Linear Algebra Operations

def dot(a, b):
    """
    Dot product of two sparse arrays.
    
    For 1-D arrays, computes inner product. For 2-D arrays, computes matrix 
    multiplication. For higher dimensions, sums products over last axis of a 
    and second-to-last axis of b.
    
    Parameters:
    - a: sparse array, first input
    - b: sparse array, second input
    
    Returns:
    Sparse array result of dot product
    """

def matmul(x1, x2):
    """
    Matrix multiplication of sparse arrays.
    
    Implements the @ operator for sparse matrices. Follows NumPy broadcasting 
    rules for batch matrix multiplication.
    
    Parameters:
    - x1: sparse array, first matrix operand
    - x2: sparse array, second matrix operand
    
    Returns:
    Sparse array result of matrix multiplication
    """

def vecdot(x1, x2, axis=-1):
    """
    Vector dot product along specified axis.
    
    Computes sum of element-wise products along specified axis, treating 
    input arrays as collections of vectors.
    
    Parameters:
    - x1: sparse array, first vector array
    - x2: sparse array, second vector array  
    - axis: int, axis along which to compute dot product
    
    Returns:
    Sparse array with dot products along specified axis
    """

def outer(a, b):
    """
    Outer product of two sparse arrays.
    
    Computes outer product of flattened input arrays, resulting in a 
    matrix where result[i,j] = a.flat[i] * b.flat[j].
    
    Parameters:
    - a: sparse array, first input vector
    - b: sparse array, second input vector
    
    Returns:
    Sparse 2-D array with outer product
    """

def kron(a, b):
    """
    Kronecker product of two sparse arrays.
    
    Computes Kronecker product, also known as tensor product. Result has 
    shape (a.shape[0]*b.shape[0], a.shape[1]*b.shape[1], ...).
    
    Parameters:
    - a: sparse array, first input
    - b: sparse array, second input
    
    Returns:
    Sparse array with Kronecker product
    """

def tensordot(a, b, axes=2):
    """
    Tensor dot product along specified axes.
    
    Computes tensor contraction by summing products over specified axes.
    More general than matrix multiplication.
    
    Parameters:
    - a: sparse array, first tensor
    - b: sparse array, second tensor
    - axes: int or sequence, axes to contract over
           - int: contract over last N axes of a and first N axes of b
           - sequence: explicit axis pairs to contract
    
    Returns:
    Sparse array result of tensor contraction
    """

def einsum(subscripts, *operands):
    """
    Einstein summation over sparse arrays.
    
    Computes tensor contractions using Einstein notation. Provides flexible
    way to specify multi-dimensional array operations.
    
    Parameters:
    - subscripts: str, Einstein summation subscripts (e.g., 'ij,jk->ik')
    - operands: sparse arrays, input tensors for the operation
    
    Returns:
    Sparse array result of Einstein summation
    """

def matrix_transpose(x):
    """
    Transpose the last two dimensions of sparse array.
    
    For matrices (2-D), equivalent to standard transpose. For higher 
    dimensional arrays, transposes only the last two axes.
    
    Parameters:
    - x: sparse array, input with at least 2 dimensions
    
    Returns:
    Sparse array with last two dimensions transposed
    """

import sparse
import numpy as np

# Create sparse matrices
A = sparse.COO.from_numpy(np.array([[1, 0, 2], [0, 3, 0], [4, 0, 5]]))
B = sparse.COO.from_numpy(np.array([[2, 1], [0, 1], [1, 0]]))

# Matrix multiplication
C = sparse.matmul(A, B)           # Matrix product A @ B
C_alt = sparse.dot(A, B)          # Equivalent using dot

print(f"A shape: {A.shape}, B shape: {B.shape}")
print(f"Result shape: {C.shape}")
print(f"Result nnz: {C.nnz}")

# Vector dot products
v1 = sparse.COO.from_numpy(np.array([1, 0, 3, 0, 2]))
v2 = sparse.COO.from_numpy(np.array([2, 1, 0, 1, 0]))

# Inner product of vectors
inner_prod = sparse.dot(v1, v2)
print(f"Inner product: {inner_prod.todense()}")  # Scalar result

# Outer product of vectors  
outer_prod = sparse.outer(v1, v2)
print(f"Outer product shape: {outer_prod.shape}")
print(f"Outer product nnz: {outer_prod.nnz}")

# Batch matrix multiplication with 3D arrays
batch_A = sparse.random((5, 10, 20), density=0.1)  # 5 matrices of 10x20
batch_B = sparse.random((5, 20, 15), density=0.1)  # 5 matrices of 20x15

# Batch matrix multiplication
batch_C = sparse.matmul(batch_A, batch_B)  # Result: 5 matrices of 10x15
print(f"Batch result shape: {batch_C.shape}")

# Vector dot product along specific axis
vectors = sparse.random((100, 50), density=0.05)    # 100 vectors of length 50
weights = sparse.ones((50,))                         # Weight vector

# Compute weighted sum for each vector
weighted_sums = sparse.vecdot(vectors, weights, axis=1)
print(f"Weighted sums shape: {weighted_sums.shape}")

# Kronecker product
A_small = sparse.COO.from_numpy(np.array([[1, 2], [3, 0]]))
B_small = sparse.COO.from_numpy(np.array([[0, 1], [1, 1]]))

kron_prod = sparse.kron(A_small, B_small)
print(f"Kronecker product shape: {kron_prod.shape}")  # (4, 4)
print(f"Kronecker product nnz: {kron_prod.nnz}")

# Tensor dot product with different contraction modes
tensor_A = sparse.random((3, 4, 5), density=0.2)
tensor_B = sparse.random((5, 6, 7), density=0.2)

# Contract over one axis (default axes=1)
result_1 = sparse.tensordot(tensor_A, tensor_B, axes=1)  # Shape: (3, 4, 6, 7)

# Contract over specific axes
result_2 = sparse.tensordot(tensor_A, tensor_B, axes=([2], [0]))  # Same as above

print(f"Tensor contraction shape: {result_1.shape}")

# Matrix multiplication using einsum
A = sparse.random((50, 30), density=0.1)
B = sparse.random((30, 40), density=0.1)

# Matrix multiply: 'ij,jk->ik'
C_einsum = sparse.einsum('ij,jk->ik', A, B)
C_matmul = sparse.matmul(A, B)

print(f"Results are equivalent: {np.allclose(C_einsum.todense(), C_matmul.todense())}")

# Batch inner product: 'bi,bi->b'
batch_vectors = sparse.random((10, 100), density=0.05)
inner_products = sparse.einsum('bi,bi->b', batch_vectors, batch_vectors)
print(f"Batch inner products shape: {inner_products.shape}")

# Trace of matrices: 'ii->'
square_matrix = sparse.random((20, 20), density=0.1)
trace = sparse.einsum('ii->', square_matrix)
print(f"Matrix trace: {trace.todense()}")

# Matrix transpose operations
matrix_3d = sparse.random((5, 10, 15), density=0.1)

# Standard transpose (swap all axes)
full_transpose = matrix_3d.transpose()
print(f"Full transpose shape: {full_transpose.shape}")  # (15, 10, 5)

# Matrix transpose (last two axes only)
matrix_transpose = sparse.matrix_transpose(matrix_3d)
print(f"Matrix transpose shape: {matrix_transpose.shape}")  # (5, 15, 10)

# Manual axis specification
custom_transpose = matrix_3d.transpose((0, 2, 1))  # Same as matrix_transpose
print(f"Custom transpose shape: {custom_transpose.shape}")  # (5, 15, 10)

# Operations between sparse and dense arrays
sparse_matrix = sparse.random((100, 50), density=0.05)
dense_vector = np.random.randn(50)

# Matrix-vector product (broadcasting)
result = sparse.dot(sparse_matrix, dense_vector)
print(f"Matrix-vector result shape: {result.shape}")
print(f"Result is sparse: {isinstance(result, sparse.SparseArray)}")

# Mixed tensor operations
dense_tensor = np.random.randn(3, 50, 4)
sparse_tensor = sparse.random((4, 10), density=0.1)

mixed_result = sparse.tensordot(dense_tensor, sparse_tensor, axes=([2], [0]))
print(f"Mixed tensor result shape: {mixed_result.shape}")  # (3, 50, 10)

# Convert to GCXS for repeated matrix operations
sparse_csr = sparse_matrix.asformat('gcxs')
for vector in vector_batch:
    result = sparse.dot(sparse_csr, vector)  # More efficient

# Use appropriate einsum subscripts for clarity and optimization
# 'ij,jk->ik' is optimized as matrix multiplication
# 'ij,ik->jk' may be less efficient than transpose + matmul