jax.ref.AbstractRef#

class jax.ref.AbstractRef(inner_aval, memory_space=None, kind=None)[source]#

Abstract mutable array reference.

Refer to the Ref guide for more information.

Parameters:

  • inner_aval (core.AbstractValue)

  • memory_space (Any)

  • kind (Any)

__init__(inner_aval, memory_space=None, kind=None)[source]#
Parameters:

  • inner_aval (core.AbstractValue)

  • memory_space (Any)

  • kind (Any)

Methods

__init__(inner_aval[, memory_space, kind])

at_least_vspace()

lo_ty()

lo_ty_qdd(qdd)

lower_val(ref)

normalize()

raise_val(*vals)

str_short([short_dtypes, mesh_axis_types])

strip_weak_type()

to_cotangent_aval()

to_tangent_aval()

update([inner_aval, memory_space, kind])

update_vma(vma)

update_weak_type(weak_type)

Attributes

inner_aval

memory_space

kind

T

addupdate

at

bitcast

dtype

get

has_qdd

is_high

bool(x) -> bool

ndim

reshape

set

shape

sharding

size

swap

transpose

vma

weak_type