jax.ShapeDtypeStruct

class jax.ShapeDtypeStruct(shape, dtype, *, sharding=None, weak_type=False, manual_axis_type=None, is_ref=False)

用于存储数组形状 (shape)、数据类型 (dtype) 及其他静态属性的容器。

ShapeDtypeStruct 通常与 jax.eval_shape() 结合使用。

参数:

• shape (Any) – 表示数组形状的整数序列

• dtype (Any) – 类 dtype 对象

• sharding – (可选) jax.Sharding 对象

• weak_type (Any)

• manual_axis_type (Any)

• is_ref (Any)

__init__(shape, dtype, *, sharding=None, weak_type=False, manual_axis_type=None, is_ref=False)

方法

__init__(shape, dtype, *[, sharding, ...])
update(**kwargs)

属性

shape
dtype
weak_type
manual_axis_type
is_ref
format
ndim
sharding
size