Skip to content

mimo_transfer_mat is jax/autograd illegal #427

Description

@rgcoe

@cmichelenstrofer Discovered that the following line in the mimo_transfer_mat function is jax/autograd illegal (in-place assignment)!

elem[idof][jdof] = block_diag(*blocks)

However, based on some further reading of the jax documentation, they (at some point) have added a workaround for this where you do something like the following.

import jax.numpy as jnp

jax_array = jnp.zeros((3,3), dtype=jnp.float32)
# jax_array[1, :] = 1.0 # can't do this
updated_array = jax_array.at[1, :].set(1.0) # can do this!

There also appears to be better error tracking/handling in jax than autograd:

In [5]: jax_array[1,1] = 1.0
TypeError: JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html

We need to either rewrite the offending line or perhaps revist using jax (#118).

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    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