Skip to content

Implement JAX-style array updates #906

Description

@tomwhite

Cubed arrays are immutable so we could implement updates in the same way as JAX, so instead of x[ind] = y, we have

x.at[idx].set(y)

According to https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html#array-updates-x-at-idx-set-y:

However, inside jit-compiled code, if the input value x of x.at[idx].set(y) is not reused, the compiler will optimize the array update to occur in-place.

It would be good if we could do a similar optimisation Cubed.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions