Tensor Abstractions

Overview

Tensor libraries define different abstractions for multi-dimensional arrays:

• NumPy: CPU, mutable, static dtype
• PyTorch: CPU/GPU, mutable, autograd
• JAX: GPU/TPU, immutable, functional, traced
• CuPy: GPU variant of NumPy
• TensorFlow: Different semantics from PyTorch

Understanding these differences is critical for building ML systems that work across frameworks.

---

The Tensor Protocol: Array-Like Objects

NumPy Array Interface

All major frameworks implement the NumPy array interface:

import numpy as np
import torch
import jax.numpy as jnp

# All implement __array__ protocol
x_np = np.array([1, 2, 3])
x_torch = torch.tensor([1, 2, 3])
x_jax = jnp.array([1, 2, 3])

# Convert to NumPy
print(np.asarray(x_torch)) # Calls x_torch.__array__()
print(np.asarray(x_jax)) # Calls x_jax.__array__()

__array_interface__ Protocol

Low-level protocol for zero-copy sharing:

class MyTensor:
 def __init__(self, data):
 self.data = np.array(data)

 @property
 def __array_interface__(self):
 """Enable zero-copy access to underlying buffer."""
 return {
 'shape': self.data.shape,
 'typestr': self.data.dtype.str,
 'data': (self.data.data, False),
 'version': 3,
 }

# Other libraries can access data without copying
t = MyTensor([1, 2, 3])
arr = np.asarray(t) # Zero-copy if possible

__array_ufunc__ Protocol

Allows frameworks to override NumPy operations:

class CustomTensor:
 def __init__(self, data):
 self.data = np.array(data)

 def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
 """Called when NumPy ufunc operates on this object."""
 if ufunc == np.add:
 # Custom addition logic
 return CustomTensor(self.data + inputs[1].data)
 return NotImplemented

# Usage
t1 = CustomTensor([1, 2])
t2 = CustomTensor([3, 4])
result = np.add(t1, t2) # Calls t1.__array_ufunc__

---

PyTorch Tensors

Fundamental Characteristics

import torch

# Creation
x = torch.tensor([1, 2, 3]) # From Python list
x = torch.zeros(3, 4) # Zeros
x = torch.randn(3, 4) # Random normal
x = torch.arange(10) # Range
x = torch.linspace(0, 1, 100) # Linspace

# Device placement
x = x.to('cuda') # Move to GPU
x = x.to(torch.device('cuda:0')) # Explicit device
x = x.cuda() # Shorthand
x = x.cpu() # Move to CPU

# Data type
x = x.float() # torch.float32
x = x.double() # torch.float64
x = x.half() # torch.float16
x = x.int() # torch.int32
x = torch.tensor([1, 2], dtype=torch.float32)

Mutability and Gradients

# PyTorch tensors are mutable
x = torch.tensor([1.0, 2.0, 3.0])
x[0] = 5.0 # In-place modification
print(x) # tensor([5., 2., 3.])

# Requires grad for autograd
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2
loss = y.sum()
loss.backward()
print(x.grad) # tensor([2., 4.])

# In-place ops affect gradient computation
x = torch.tensor([1.0, 2.0], requires_grad=True)
x += 1 # In-place add
# Later operations will see modified x

Memory Layout

# Row-major (C-contiguous) vs Column-major (Fortran-contiguous)
x = torch.randn(3, 4)
print(x.is_contiguous()) # True (C-contiguous by default)

# Transposed is not contiguous
y = x.T
print(y.is_contiguous()) # False

# Some operations require contiguous
z = y.contiguous() # Makes copy if needed
print(z.is_contiguous()) # True

# Important for GPU kernels and backward pass

Indexing and Slicing

x = torch.randn(3, 4, 5)

# Basic indexing
x[0] # First row
x[0, 1] # Element [0, 1]
x[:, 1,:] # Middle column, all others

# Advanced indexing
idx = torch.tensor([0, 2])
x[idx] # Rows 0 and 2

# Fancy indexing
mask = x > 0
x[mask] # Elements > 0 (flattened)

# Indexing creates views (memory shared)
y = x[:, 0]
y[0] = 999 # Modifies x too!

---

JAX Arrays

Immutability

JAX arrays are immutable by design (enable JIT compilation):

import jax.numpy as jnp

x = jnp.array([1, 2, 3])

# Can't modify in-place
# x[0] = 5 # TypeError!

# Must use functional updates
x = x.at[0].set(5)
print(x) # [5, 2, 3] (new array)

# Batch operations
x = x.at[jnp.array([0, 2])].set(99)
print(x) # [99, 2, 99]

# Add instead of increment
x = x.at[0].add(10)
print(x) # [109, 2, 99]

Shape Tracing

JAX uses shape tracing during JIT compilation:

import jax
import jax.numpy as jnp

@jax.jit
def sum_array(x):
 return jnp.sum(x)

# Shape is traced at compile time
x = jnp.array([1, 2, 3])
result = sum_array(x) # JIT compiles with shape (3,)

# Different shape requires recompilation
y = jnp.array([1, 2, 3, 4, 5])
result = sum_array(y) # Recompiles with shape (5,)

Functional Transformations

import jax
import jax.numpy as jnp

def loss_fn(params, x, y):
 pred = jnp.dot(x, params)
 return jnp.mean((pred - y) ** 2)

# Gradient
grad_fn = jax.grad(loss_fn)
grads = grad_fn(params, x, y) # Returns gradients w.r.t. params

# Vectorize
batched_loss = jax.vmap(loss_fn, in_axes=(None, 0, 0))
losses = batched_loss(params, x_batch, y_batch) # Vectorized

# Compose transformations
@jax.jit
@jax.vmap
def batched_forward(params, batch_x):
 return jnp.dot(batch_x, params)

-

NumPy Arrays

Core Characteristics

import numpy as np

# Creation
x = np.array([1, 2, 3])
x = np.zeros((3, 4))
x = np.random.randn(3, 4)

# CPU only
x = x # No device placement

# Mutable
x[0] = 5 # In-place modification

# No autograd
# No backward() method

# Broadcasting
x = np.array([1, 2, 3](/1,-2,-3/)) # Shape (1, 3)
y = np.array([[1], [2], [3]]) # Shape (3, 1)
result = x + y # Broadcasts to (3, 3)

Broadcasting Rules

Critical for ML code:

import numpy as np

# Rule 1
# (3, 1, 4) broadcasts with (1, 3, 4)
# Result

x = np.zeros((3, 1, 4))
y = np.zeros((1, 3, 4))
result = x + y
print(result.shape) # (3, 3, 4)

# Rule 2
# (3, 4) broadcasts with (4,)
x = np.zeros((3, 4))
y = np.zeros(4)
result = x + y # y broadcasted to (1, 4) then (3, 4)
print(result.shape) # (3, 4)

# Practical example
batch_size = 32
seq_len = 10
vocab_size = 50000

logits = np.random.randn(batch_size, seq_len, vocab_size)
temperature = 0.5 # Scalar

scaled_logits = logits / temperature # Broadcasting works
print(scaled_logits.shape) # (32, 10, 50000)

-

Practical: Cross-Framework Code

Write Once, Run Anywhere

from typing import Protocol, Union
import numpy as np
import torch
import jax.numpy as jnp

class ArrayLike(Protocol):
 """Any array-like object."""
 shape: tuple
 dtype: object

 def __array__(self) -> np.ndarray:
...

def normalize(arr: ArrayLike) -> ArrayLike:
 """Works with NumPy, PyTorch, JAX arrays."""
 mean = np.mean(arr)
 std = np.std(arr)
 return (arr - mean) / std

# All work!
normalize(np.array([1, 2, 3]))
normalize(torch.tensor([1, 2, 3]))
normalize(jnp.array([1, 2, 3]))

Framework-Agnostic Training Loop

import numpy as np
import torch

def train_step(model, x, y, optimizer, framework='torch'):
 """Training step works with PyTorch or NumPy models."""

 if framework == 'torch':
 # PyTorch path
 optimizer.zero_grad()
 pred = model(x)
 loss = torch.nn.functional.mse_loss(pred, y)
 loss.backward()
 optimizer.step()
 return float(loss)

 elif framework == 'numpy':
 # NumPy path (manual backprop)
 pred = model.forward(x)
 loss = np.mean((pred - y) ** 2)
 grad = 2 * (pred - y) / len(y)
 model.backward(grad)
 model.update()
 return float(loss)

Conversion Utilities

import numpy as np
import torch
import jax.numpy as jnp

def to_numpy(x):
 """Convert any array to NumPy."""
 if isinstance(x, np.ndarray):
 return x
 elif isinstance(x, torch.Tensor):
 return x.detach().cpu().numpy()
 elif isinstance(x, jnp.ndarray):
 return np.array(x)
 else:
 return np.asarray(x)

def to_torch(x, device='cpu'):
 """Convert any array to PyTorch."""
 if isinstance(x, torch.Tensor):
 return x.to(device)
 else:
 x_np = to_numpy(x)
 return torch.from_numpy(x_np).to(device)

def to_jax(x):
 """Convert any array to JAX."""
 if isinstance(x, jnp.ndarray):
 return x
 else:
 x_np = to_numpy(x)
 return jnp.array(x_np)

# Usage
x_np = np.array([1, 2, 3])
x_torch = to_torch(x_np, device='cuda')
x_jax = to_jax(x_torch)
x_np = to_numpy(x_jax)

-

Tensor Shape and Dimension Operations

Common Operations

import torch

x = torch.randn(3, 4, 5)

# Shape inspection
print(x.shape) # torch.Size([3, 4, 5])
print(x.ndim) # 3
print(x.size()) # torch.Size([3, 4, 5])

# Reshape
y = x.reshape(12, 5) # (3, 4, 5) → (12, 5)
y = x.view(12, 5) # Requires contiguous (faster)
y = x.flatten() # Flatten to 1D

# Transpose
y = x.T # Full transpose
y = x.transpose(0, 2) # Swap dims 0 and 2

# Squeeze/Unsqueeze
x_sq = torch.randn(1, 3, 1, 4)
y = x_sq.squeeze() # Remove dims of size 1: (3, 4)
y = x.unsqueeze(0) # Add dimension at 0: (1, 3, 4, 5)

# Permute
y = x.permute(2, 0, 1) # Reorder: (5, 3, 4)

# Expand
y = x.expand(3, 4, 10) # Only works if target size ≥ original

# Repeat
y = x.repeat(2, 1, 1) # Copy tensor: (6, 4, 5)

Broadcasting vs Reshaping

import torch

# Broadcasting
x = torch.zeros(3, 1, 4)
y = torch.zeros(1, 3, 4)
result = x + y # Broadcasts to (3, 3, 4) without copying

# Reshaping
z = torch.randn(12)
w = z.reshape(3, 4) # May or may not copy depending on layout

# Difference in inference:
# - Broadcasting
# - Reshaping

---

Device Management Patterns

Automatic Device Handling

import torch

class Model(torch.nn.Module):
 def __init__(self):
 super().__init__()
 self.linear = torch.nn.Linear(10, 5)

 @property
 def device(self):
 """Infer device from parameters."""
 return next(self.parameters()).device

 def to_device(self, x):
 """Move input to same device as model."""
 return x.to(self.device)

 def forward(self, x):
 x = self.to_device(x)
 return self.linear(x)

# Usage
model = Model()
model.to('cuda')

# Forward works even if input is on CPU
x_cpu = torch.randn(3, 10)
output = model(x_cpu) # Automatically moved to GPU

Mixed Precision with Automatic Device Placement

import torch

# Model on GPU, but some computations in float32
model = torch.nn.Linear(10, 5).to('cuda').half()

# Input automatically cast to float16
x = torch.randn(3, 10, device='cuda', dtype=torch.float16)
y = model(x) # Computed in float16

# Some operations require float32
with torch.autocast(device_type='cuda', dtype=torch.float32):
 # All operations in float32
 z = torch.nn.functional.softmax(y, dim=-1)

-

Practical: Implementing Tensor-Agnostic Code

from typing import Union, Optional
import numpy as np
import torch

Array = Union[np.ndarray, torch.Tensor]

class TensorProcessor:
 """Works with NumPy or PyTorch tensors."""

 @staticmethod
 def is_torch(x):
 return isinstance(x, torch.Tensor)

 @staticmethod
 def zero_like(x: Array) -> Array:
 """Create zero tensor same shape/dtype as input."""
 if TensorProcessor.is_torch(x):
 return torch.zeros_like(x)
 else:
 return np.zeros_like(x)

 @staticmethod
 def concatenate(arrays: list[Array], axis: int = 0) -> Array:
 """Concatenate along axis."""
 if TensorProcessor.is_torch(arrays[0]):
 return torch.cat(arrays, dim=axis)
 else:
 return np.concatenate(arrays, axis=axis)

 @staticmethod
 def normalize(x: Array) -> Array:
 """Normalize to [0, 1]."""
 x_min = x.min()
 x_max = x.max()
 return (x - x_min) / (x_max - x_min + 1e-8)

# Usage
x_np = np.array([1, 2, 3])
x_torch = torch.tensor([1, 2, 3])

norm_np = TensorProcessor.normalize(x_np)
norm_torch = TensorProcessor.normalize(x_torch)

---

Summary

Feature	NumPy	PyTorch	JAX
Mutability	Mutable	Mutable	Immutable
Devices	CPU only	CPU/GPU/TPU	GPU/TPU
Autograd	No	Yes	Via transformations
JIT	No	TorchScript	Native
Broadcasting	Yes	Yes	Yes
Functional	No	No	Yes
Array Interface	Yes	Yes	Yes
Use Case	General	Production ML	Research/JAX

-

Related Topics

• 02 Autograd Implementation - How PyTorch tracks gradients
• 05 Custom Operators - Extending tensors
• 06 Inference Optimization Patterns - Efficient tensor operations
• 00 Readme - Tensor implementation in C/C++