Issue 1:
Description

When multiplying two sparse BCOO matrices it seems the result always stores explicit zero-entries even when the corresponding row/column of a and b are all zero:

import jax
import numpy as np
a = jax.experimental.sparse.BCOO.fromdense(np.diag([1., 2.]))
b = jax.experimental.sparse.BCOO.fromdense(np.diag([3., 4.]))
(a @ b).data, (a @ b).indices
>>> (Array([3., 0., 0., 8.], dtype=float64),
Array([[0, 0],
[0, 1],
[1, 0],
[1, 1]], dtype=int32))

Expected output:

>>> (Array([3., 8.], dtype=float64),
Array([[0, 0],
[1, 1]], dtype=int32))

What jax/jaxlib version are you using?

0.4.8

Which accelerator(s) are you using?

GPU

Additional system info

No response

NVIDIA GPU info

No response

Issue 2:
Please:

Check for duplicate issues.

Does jaxlib support building with GCC 11? I'm seeing a number of errors.

If applicable, include full error messages/tracebacks.

ERROR: /build/output/external/org_tensorflow/tensorflow/tools/proto_text/BUILD:31:10: Linking external/org_tensorflow/tensorflow/tools/proto_text/gen_proto_text_functions failed: (Exit 1): crosstool_wrapper_driver_is_not_gcc failed: error executing command external/local_config_cuda/crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc @bazel-out/k8-opt/bin/external/org_tensorflow/tensorflow/tools/proto_text/gen_proto_text_functions-2.params
/nix/store/m2vh2ny7bqpwij1gpmvl5gxj7y4dgr4f-binutils-2.38/bin/ld: /nix/store/x0gmnlsh4mzfsngvxj9ahsc1vm5sdrg5-protobuf-3.19.3/lib/libprotobuf.a(parser.o): in function `std::_Hashtable<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const, google::protobuf::FieldDescriptorProto_Type>, std::allocator<std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const, google::protobuf::FieldDescriptorProto_Type> >, std::__detail::_Select1st, std::equal_to<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::hash<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::__detail::_Mod_range_hashing, std::__detail::_Default_ranged_hash, std::__detail::_Prime_rehash_policy, std::__detail::_Hashtable_traits<true, false, true> >::_M_rehash(unsigned long, unsigned long const&)':
(.text._ZNSt10_HashtableINSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEESt4pairIKS5_N6google8protobuf25FieldDescriptorProto_TypeEESaISB_ENSt8__detail10_Select1stESt8equal_toIS5_ESt4hashIS5_ENSD_18_Mod_range_hashingENSD_20_Default_ranged_hashENSD_20_Prime_rehash_policyENSD_17_Hashtable_traitsILb1ELb0ELb1EEEE9_M_rehashEmRKm[_ZNSt10_HashtableINSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEESt4pairIKS5_N6google8protobuf25FieldDescriptorProto_TypeEESaISB_ENSt8__detail10_Select1stESt8equal_toIS5_ESt4hashIS5_ENSD_18_Mod_range_hashingENSD_20_Default_ranged_hashENSD_20_Prime_rehash_policyENSD_17_Hashtable_traitsILb1ELb0ELb1EEEE9_M_rehashEmRKm]+0xff): undefined reference to `std::__throw_bad_array_new_length()'
/nix/store/m2vh2ny7bqpwij1gpmvl5gxj7y4dgr4f-binutils-2.38/bin/ld: /nix/store/x0gmnlsh4mzfsngvxj9ahsc1vm5sdrg5-protobuf-3.19.3/lib/libprotobuf.a(descriptor.o): in function `bool google::protobuf::InsertIfNotPresent<std::unordered_map<std::pair<void const*, google::protobuf::stringpiece_internal::StringPiece>, google::protobuf::FieldDescriptor const*, google::protobuf::(anonymous namespace)::PointerStringPairHash, std::equal_to<std::pair<void const*, google::protobuf::stringpiece_internal::StringPiece> >, std::allocator<std::pair<std::pair<void const*, google::protobuf::stringpiece_internal::StringPiece> const, google::protobuf::FieldDescriptor const*> > > >(std::unordered_map<std::pair<void const*, google::protobuf::stringpiece_internal::StringPiece>, google::protobuf::FieldDescriptor const*, google::protobuf::(anonymous namespace)::PointerStringPairHash, std::equal_to<std::pair<void const*, google::protobuf::stringpiece_internal::StringPiece> >, std::allocator<std::pair<std::pair<void const*, google::protobuf::stringpiece_internal::StringPiece> const, google::protobuf::FieldDescriptor const*> > >*, std::unordered_map<std::pair<void const*, google::protobuf::stringpiece_internal::StringPiece>, google::protobuf::FieldDescriptor const*, google::protobuf::(anonymous namespace)::PointerStringPairHash, std::equal_to<std::pair<void const*, google::protobuf::stringpiece_internal::StringPiece> >, std::allocator<std::pair<std::pair<void const*, google::protobuf::stringpiece_internal::StringPiece> const, google::protobuf::FieldDescriptor const*> > >::value_type::first_type const&, std::unordered_map<std::pair<void const*, google::protobuf::stringpiece_internal::StringPiece>, google::protobuf::FieldDescriptor const*, google::protobuf::(anonymous namespace)::PointerStringPairHash, std::equal_to<std::pair<void const*, google::protobuf::stringpiece_internal::StringPiece> >, std::allocator<std::pair<std::pair<void const*, google::protobuf::stringpiece_internal::StringPiece> const, google::protobuf::FieldDescriptor const*> > >::value_type::second_type const&) [clone .isra.0]':
(.text+0xd33): undefined reference to `std::__throw_bad_array_new_length()'
/nix/store/m2vh2ny7bqpwij1gpmvl5gxj7y4dgr4f-binutils-2.38/bin/ld: /nix/store/x0gmnlsh4mzfsngvxj9ahsc1vm5sdrg5-protobuf-3.19.3/lib/libprotobuf.a(descriptor.o): in function `std::pair<std::__detail::_Node_iterator<google::protobuf::Symbol, true, true>, bool> std::_Hashtable<google::protobuf::Symbol, google::protobuf::Symbol, std::allocator<google::protobuf::Symbol>, std::__detail::_Identity, google::protobuf::(anonymous namespace)::FieldsByNumberEq, google::protobuf::(anonymous namespace)::FieldsByNumberHash, std::__detail::_Mod_range_hashing, std::__detail::_Default_ranged_hash, std::__detail::_Prime_rehash_policy, std::__detail::_Hashtable_traits<true, true, true> >::_M_insert<google::protobuf::Symbol, std::__detail::_AllocNode<std::allocator<std::__detail::_Hash_node<google::protobuf::Symbol, true> > > >(google::protobuf::Symbol&&, std::__detail::_AllocNode<std::allocator<std::__detail::_Hash_node<google::protobuf::Symbol, true> > > const&, std::integral_constant<bool, true>) [clone .constprop.0] [clone .isra.0]':
(.text+0x31bf): undefined reference to `std::__throw_bad_array_new_length()'
/nix/store/m2vh2ny7bqpwij1gpmvl5gxj7y4dgr4f-binutils-2.38/bin/ld: /nix/store/x0gmnlsh4mzfsngvxj9ahsc1vm5sdrg5-protobuf-3.19.3/lib/libprotobuf.a(descriptor.o): in function `google::protobuf::FileDescriptorTables::AddAliasUnderParent(void const*, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, google::protobuf::Symbol)':
(.text+0x4373): undefined reference to `std::__throw_bad_array_new_length()'
/nix/store/m2vh2ny7bqpwij1gpmvl5gxj7y4dgr4f-binutils-2.38/bin/ld: /nix/store/x0gmnlsh4mzfsngvxj9ahsc1vm5sdrg5-protobuf-3.19.3/lib/libprotobuf.a(descriptor.o): in function `google::protobuf::DescriptorPool::Tables::AddSymbol(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, google::protobuf::Symbol)':
(.text+0x109cf): undefined reference to `std::__throw_bad_array_new_length()'
/nix/store/m2vh2ny7bqpwij1gpmvl5gxj7y4dgr4f-binutils-2.38/bin/ld: /nix/store/x0gmnlsh4mzfsngvxj9ahsc1vm5sdrg5-protobuf-3.19.3/lib/libprotobuf.a(descriptor.o):(.text+0x1ee0f): more undefined references to `std::__throw_bad_array_new_length()' follow
collect2: error: ld returned 1 exit status
Target //build:build_wheel failed to build
Use --verbose_failures to see the command lines of failed build steps.
ERROR: /build/output/external/org_tensorflow/tensorflow/core/framework/BUILD:1400:31 Middleman _middlemen/@org_Utensorflow_S_Stensorflow_Score_Sframework_Cattr_Uvalue_Uproto_Utext-BazelCppSemantics_build_arch_k8-opt failed: (Exit 1): crosstool_wrapper_driver_is_not_gcc failed: error executing command external/local_config_cuda/crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc @bazel-out/k8-opt/bin/external/org_tensorflow/tensorflow/tools/proto_text/gen_proto_text_functions-2.params

What do I need to do in order to build with GCC 11?