fix(math): ShardedArray pytree flatten + remove_diag guard#839

chaoming0625 merged 1 commit into

masterbrainpy/BrainPy:masterfrom

fix/audit-20260619-math-corebrainpy/BrainPy:fix/audit-20260619-math-coreCopy head branch name to clipboard