jax.lax.broadcast_shapes# jax.lax.broadcast_shapes(*shapes: tuple[int, ...]) → tuple[int, ...][source]# jax.lax.broadcast_shapes(*shapes: tuple[int | core.Tracer, ...]) → tuple[int | core.Tracer, ...] 返回 NumPy 广播 shapes 后得到的形状。
