jax.ref.AbstractRef#

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

Abstract mutable array reference.

Refer to the ArrayRef guide for more information.

Parameters:

  • inner_aval (core.AbstractValue)

  • memory_space (Any)

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

  • inner_aval (core.AbstractValue)

  • memory_space (Any)

Methods

__init__(inner_aval[, memory_space])

at_least_vspace()

lo_ty()

lo_ty_qdd(qdd)

normalize()

str_short([short_dtypes, mesh_axis_types])

strip_weak_type()

to_tangent_aval()

update([inner_aval, memory_space])

update_vma(vma)

update_weak_type(weak_type)

Attributes

inner_aval

memory_space

T

addupdate

at

bitcast

dtype

get

has_qdd

is_high

ndim

reshape

set

shape

sharding

size

swap

transpose

vma

weak_type