Doesn't support tensors right now, as it's somewhat hairy to support both at the same time. Since we use a generic lowering the result is messy and needs a mem2reg pass to eliminate extra load/store/allocas. PiperOrigin-RevId: 339562971