N.B. for the case of jnp.where(), specifying x and y args other than 'None' resolves the issue
Post
Replies
Boosts
Views
Activity
and for jnp.round the error is:
jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: :1:0: error: failed to legalize operation 'mhlo.round_nearest_even'
Same here. More specifically the error is:
jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: :1:0: error: failed to legalize operation 'mhlo.pad'