JAX 中的数据类型

JAX 数组中的数据类型与 NumPy 中的类似。例如，下面展示了如何在 JAX 中创建浮点数（float）和整数（int）类型的数据。

1import jax.numpy as jnp 
2
3x = jnp.float32(1.25844)
4print("x :", x) 
5
6y = jnp.int32(45.25844)
7print("y :", y)

当我们检查这些数据的类型时，会发现它们都是 DeviceArray 类型。从以下代码可以看出，无论是 float32 还是 int32 变量，其类型都相同。

1import jax.numpy as jnp 
2
3x = jnp.float32(1.25844)
4print("type of x: ",type(x))
5
6y = jnp.int32(45.25844)
7print("type of y: ",type(y))

JAX 中的 DeviceArray 相当于 NumPy 中的 numpy.ndarray，而 jax.numpy 提供了与 NumPy 非常相似的接口。不过，JAX 还提供了 jax.lax，这是一个功能更强大但要求更严格的底层 API。例如，在 jax.numpy 中可以对不同类型的数据进行加法运算，但 jax.lax 不允许这种混合类型操作。

创建 JAX 数组的方法

我们可以像在 NumPy 中一样创建 JAX 数组。例如，可以使用以下方法：

• arange() 函数
• linspace() 函数
• Python 列表
• zeros() 函数
• ones() 函数
• identity() 或 eye() 函数

下面我们来看上述函数的实际输出结果：

1import jax.numpy as jnp 
2
3a = jnp.arange(10)
4print("a : ", a)
5
6b = jnp.linspace(0, 10, 30)
7print("b :", b)
8
9scores = [50,60,70,30,25,70]
10scores_array = jnp.array(scores)
11print("scores_array :", scores_array)
12
13c = jnp.zeros(5)
14print("c :", c)
15
16d = jnp.ones(5)
17print("d :", d)
18
19e = jnp.eye(5)
20print("e :", e)
21
22f = jnp.identity(5)
23print("f :", f)