JAX中的多维数组类型

jnp.ndarray 是 JAX 中的多维数组数据类型,其接口与 NumPy 中的 numpy.ndarray 类型非常相似。

1. 创建JAX数组

import jax.numpy as jnp
import jax


print(jnp.ndarray) #  <class 'jax.Array'>
print(jnp.array) # <class 'jaxlib.xla_extension.ArrayImpl'>

## 1.1 从 Python 列表创建
# 一维数组
arr1 = jnp.array([1, 2, 3])
# 二维数组
arr2 = jnp.array([[1, 2, 3], [4, 5, 6]])
# 三维数组
arr3 = jnp.array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]])

## 1.2 创建特定值数组
# 创建一个形状为 (3, 3) 的数组,元素都是 0
zeros_array = jnp.zeros((3, 3))
# 创建一个形状为 (2, 2) 的数组,元素都是 1
ones_array = jnp.ones((2, 2))
# 创建一个形状为 (3, 3) 的数组,元素都是指定的值 5
constant_array = jnp.full((3, 3), 5)

## 1.3 创建随机数组
# 在较新版本的 JAX 中,随机数生成模块已经从 jnp 模块中分离出来,成为一个独立的 jax.random 模块。

# 生成形状为 (2, 3) 的标准正态分布随机数组
key = jax.random.PRNGKey(42)
random_normal_array = jax.random.normal(key, shape=(2, 3))


# 生成形状为 (3, 3) 的在 [0, 1) 范围内均匀分布的随机数组
key = jax.random.PRNGKey(42)
random_uniform_array = jax.random.uniform(key, shape=(3, 3))

## 1.4 使用序列生成数组:
# 创建一个从 0 到 4 步长为 1 的数组
seq_array = jnp.arange(5)

# 创建一个从 0 到 1 之间等间距的 5 个元素的数组
linspace_array = jnp.linspace(0, 1, 5)
#print(linspace_array)

2. jnp.ndarray属性

arr = jnp.array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]], [[13, 14, 15], [16, 17, 18]]])
## 2.1 shape: 返回数组的维度信息,即各个轴的长度。

shape = arr.shape
## 2.2 dtype: 返回数组中元素的数据类型。

dtype = arr.dtype

## 2.3 size: 返回数组中元素的总数。
size = arr.size

## 2. 4 ndim: 返回数组的维度数。
ndim = arr.ndim

3. jnp.ndarray方法

  1. 数学操作:

    • jnp.add, jnp.subtract, jnp.multiply, jnp.divide: 加法、减法、乘法、除法。
    • jnp.sum, jnp.mean, jnp.max, jnp.min: 求和、均值、最大值、最小值。
  2. 数组操作:

    • jnp.reshape(new_shape): 返回具有新形状的数组。
    • jnp.transpose(axes): 返回轴重新排列的数组。
    • jnp.concatenate(arrays, axis): 沿指定轴连接数组。
  3. 逻辑操作:

    • jnp.logical_and, jnp.logical_or, jnp.logical_not: 逻辑与、逻辑或、逻辑非。
  4. 线性代数:

    • jnp.dot(a, b): 点积。
    • jnp.linalg.norm(x): 范数。
  5. 索引和切片:

    • arr[index]: 获取索引处的元素。
    • arr[start:stop:step]: 对数组进行切片。
  6. 其他:

    • jnp.where(condition, x, y): 根据条件返回元素。
    • jnp.concatenate(arrays, axis): 沿指定轴连接数组。
    • jnp.argmax(arr), jnp.argmin(arr): 最大值和最小值的索引。

参考:

https://jax.readthedocs.io/en/latest/jax.numpy.html

  • 10
    点赞
  • 7
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值