968 lines
36 KiB
Python
968 lines
36 KiB
Python
import torch
|
|
import math
|
|
import torch.nn.functional as F
|
|
from .MathExprVisitor import MathExprVisitor
|
|
from ..helper_functions import generate_dim_variables
|
|
|
|
|
|
class UnifiedMathVisitor(MathExprVisitor):
|
|
def __init__(self, variables, shape=None, device=None):
|
|
self.variables = variables
|
|
self.spatial_variables = variables.copy()
|
|
self.shape = shape if shape is not None else ()
|
|
if device is None:
|
|
self.device = next((v.device for v in variables.values() if isinstance(v, torch.Tensor)), torch.device("cpu"))
|
|
else:
|
|
self.device = device
|
|
|
|
def _is_tensor(self, val):
|
|
return isinstance(val, torch.Tensor)
|
|
|
|
def _is_list(self, val):
|
|
return isinstance(val, (list, tuple))
|
|
|
|
def _promote_to_tensor(self, val):
|
|
if self._is_tensor(val):
|
|
return val.contiguous()
|
|
if self._is_list(val):
|
|
return torch.tensor(val, device=self.device)
|
|
return torch.tensor(val, device=self.device)
|
|
|
|
def _bin_op(self, a, b, torch_op, scalar_op):
|
|
"""
|
|
Generic binary operation handler.
|
|
"""
|
|
# one of them is a list and one is tensor
|
|
if self._is_tensor(a) and self._is_list(b):
|
|
if(a.shape[0]==len(b)):
|
|
c = torch.split(a,1)
|
|
return torch.cat([self._bin_op(x, y, torch_op, scalar_op) for x,y in zip(a,c)],dim=0)
|
|
return torch.cat([self._bin_op(a, x, torch_op, scalar_op) for x in b], dim=0)
|
|
if self._is_list(a) and self._is_tensor(b):
|
|
if(b.shape[0]==len(a)):
|
|
c = torch.split(a,1)
|
|
return torch.cat([self._bin_op(x, y, torch_op, scalar_op) for x,y in zip(c,b)],dim=0)
|
|
return torch.cat([self._bin_op(x, b, torch_op, scalar_op) for x in a], dim=0)
|
|
|
|
if self._is_list(a) and not self._is_tensor(b):
|
|
if self._is_list(b):
|
|
if len(a) != len(b):
|
|
raise ValueError("List length mismatch")
|
|
return [self._bin_op(x, y, torch_op, scalar_op) for x, y in zip(a, b)]
|
|
return [self._bin_op(x, b, torch_op, scalar_op) for x in a]
|
|
|
|
if not self._is_tensor(a) and self._is_list(b):
|
|
return [self._bin_op(a, x, torch_op, scalar_op) for x in b]
|
|
|
|
if self._is_tensor(a) or self._is_tensor(b):
|
|
if torch_op:
|
|
return torch_op(a, b).contiguous()
|
|
return scalar_op(a, b)
|
|
|
|
return scalar_op(a, b)
|
|
|
|
def _unary_op(self, a, torch_op, scalar_op):
|
|
if self._is_list(a):
|
|
return [self._unary_op(x, torch_op, scalar_op) for x in a]
|
|
if self._is_tensor(a):
|
|
return torch_op(a).contiguous() if torch_op else scalar_op(a)
|
|
return scalar_op(a)
|
|
|
|
def _reduction_op(self, val, torch_op, list_op):
|
|
if self._is_tensor(val):
|
|
return torch_op(val)
|
|
if self._is_list(val):
|
|
return list_op(val)
|
|
return val
|
|
|
|
# ========================
|
|
# Visitors
|
|
# ========================
|
|
|
|
def visitNumberExp(self, ctx):
|
|
val_str = ctx.getText()
|
|
if "." in val_str or "e" in val_str:
|
|
return float(val_str)
|
|
return int(val_str)
|
|
|
|
def visitConstantExp(self, ctx):
|
|
name = ctx.getText().lower()
|
|
if name == "pi":
|
|
return math.pi
|
|
if name == "e":
|
|
return math.e
|
|
raise ValueError(f"Unknown constant: {name}")
|
|
|
|
def visitVariableExp(self, ctx):
|
|
name = ctx.getText()
|
|
if name not in self.variables:
|
|
raise ValueError(f"Variable '{name}' not found")
|
|
return self.variables[name]
|
|
|
|
def visitListExp(self, ctx):
|
|
return [self.visit(e) for e in ctx.expr()]
|
|
|
|
def visitParenExp(self, ctx):
|
|
return self.visit(ctx.expr())
|
|
|
|
def visitUnaryPlus(self, ctx):
|
|
return self._unary_op(self.visit(ctx.unaryExpr()), lambda x: x, lambda x: +x)
|
|
|
|
def visitUnaryMinus(self, ctx):
|
|
return self._unary_op(self.visit(ctx.unaryExpr()), torch.neg, lambda x: -x)
|
|
|
|
# Binary Ops
|
|
def visitAddExp(self, ctx):
|
|
return self._bin_op(self.visit(ctx.addExpr()), self.visit(ctx.mulExpr()), torch.add, lambda a, b: a + b)
|
|
|
|
def visitSubExp(self, ctx):
|
|
return self._bin_op(self.visit(ctx.addExpr()), self.visit(ctx.mulExpr()), torch.sub, lambda a, b: a - b)
|
|
|
|
def visitMulExp(self, ctx):
|
|
return self._bin_op(self.visit(ctx.mulExpr()), self.visit(ctx.powExpr()), torch.mul, lambda a, b: a * b)
|
|
|
|
def visitDivExp(self, ctx):
|
|
return self._bin_op(self.visit(ctx.mulExpr()), self.visit(ctx.powExpr()), torch.div, lambda a, b: a / b)
|
|
|
|
def visitModExp(self, ctx):
|
|
return self._bin_op(self.visit(ctx.mulExpr()), self.visit(ctx.powExpr()), torch.fmod, lambda a, b: a % b)
|
|
|
|
def visitPowExp(self, ctx):
|
|
return self._bin_op(self.visit(ctx.unaryExpr()), self.visit(ctx.powExpr()), torch.pow, math.pow)
|
|
|
|
def _bool_op(self, a, b, torch_op, scalar_op):
|
|
return self._bin_op(a, b, torch_op, scalar_op)
|
|
|
|
def visitNeExp(self, ctx):
|
|
return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.ne, lambda a, b: float(a != b))
|
|
|
|
def visitEqExp(self, ctx):
|
|
return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.eq, lambda a, b: float(a == b))
|
|
|
|
def visitGtExp(self, ctx):
|
|
return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.gt, lambda a, b: float(a > b))
|
|
|
|
def visitLtExp(self, ctx):
|
|
return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.lt, lambda a, b: float(a < b))
|
|
|
|
def visitGeExp(self, ctx):
|
|
return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.ge, lambda a, b: float(a >= b))
|
|
|
|
def visitLeExp(self, ctx):
|
|
return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.le, lambda a, b: float(a <= b))
|
|
|
|
# Functions
|
|
|
|
|
|
def visitSinFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.sin, math.sin)
|
|
|
|
def visitCosFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.cos, math.cos)
|
|
|
|
def visitTanFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.tan, math.tan)
|
|
|
|
def visitAsinFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.asin, math.asin)
|
|
|
|
def visitAcosFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.acos, math.acos)
|
|
|
|
def visitAtanFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.atan, math.atan)
|
|
|
|
def visitSinhFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.sinh, math.sinh)
|
|
|
|
def visitCoshFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.cosh, math.cosh)
|
|
|
|
def visitTanhFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.tanh, math.tanh)
|
|
|
|
def visitAsinhFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.asinh, math.asinh)
|
|
|
|
def visitAcoshFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.acosh, math.acosh)
|
|
|
|
def visitAtanhFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.atanh, math.atanh)
|
|
|
|
def visitAbsFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.abs, abs)
|
|
|
|
def visitAbsExp(self, ctx):
|
|
val = self.visit(ctx.expr())
|
|
if self._is_list(val):
|
|
return torch.linalg.norm(self._promote_to_tensor(val))
|
|
if self._is_tensor(val):
|
|
return torch.linalg.norm(val)
|
|
return abs(val)
|
|
|
|
def visitSqrtFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.sqrt, math.sqrt)
|
|
|
|
def visitLnFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.log, math.log)
|
|
|
|
def visitLogFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.log10, math.log10)
|
|
|
|
def visitExpFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.exp, math.exp)
|
|
|
|
def visitFloorFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.floor, math.floor)
|
|
|
|
def visitCeilFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.ceil, math.ceil)
|
|
|
|
def visitRoundFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.round, round)
|
|
|
|
def visitSignFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.sign, lambda x: (1.0 if x > 0 else (-1.0 if x < 0 else 0.0)))
|
|
|
|
def visitFractFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), lambda x: x - torch.floor(x), lambda x: x - math.floor(x))
|
|
|
|
def visitGammaFunc(self, ctx):
|
|
torch_gamma = getattr(torch.special, "gamma", None)
|
|
if torch_gamma is None:
|
|
torch_gamma = lambda x: torch.exp(torch.lgamma(x))
|
|
return self._unary_op(self.visit(ctx.expr()), torch_gamma, math.gamma)
|
|
|
|
def visitSigmoidFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.sigmoid, lambda x: 1.0 / (1.0 + math.exp(-x)))
|
|
|
|
def visitReluFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.relu, lambda x: max(0.0, x))
|
|
|
|
def visitSoftplusFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), F.softplus, lambda x: math.log(1.0 + math.exp(x)))
|
|
|
|
def visitGeluFunc(self, ctx):
|
|
return self._unary_op(
|
|
self.visit(ctx.expr()), F.gelu, lambda x: 0.5 * x * (1 + math.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * math.pow(x, 3))))
|
|
)
|
|
|
|
def visitAnglFunc(self, ctx):
|
|
return self._unary_op(self.visit(ctx.expr()), torch.angle, lambda x: math.atan2(0, x) if x < 0 else 0)
|
|
|
|
def visitPrintFunc(self, ctx):
|
|
val = self.visit(ctx.expr())
|
|
print(f"{val}")
|
|
return val
|
|
|
|
def visitTNormFunc(self, ctx):
|
|
val = self.visit(ctx.expr())
|
|
if self._is_tensor(val):
|
|
return F.normalize(val, p=2, dim=-1)
|
|
return 1.0 if val != 0 else 0.0
|
|
|
|
def visitSNormFunc(self, ctx):
|
|
val = self.visit(ctx.expr())
|
|
if self._is_tensor(val):
|
|
return torch.linalg.norm(val)
|
|
return abs(val)
|
|
|
|
# Two-argument functions
|
|
def visitPowFunc(self, ctx):
|
|
return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), torch.pow, math.pow)
|
|
|
|
def visitAtan2Func(self, ctx):
|
|
return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), torch.atan2, math.atan2)
|
|
|
|
def visitTMinFunc(self, ctx):
|
|
return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), torch.minimum, min)
|
|
|
|
def visitTMaxFunc(self, ctx):
|
|
return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), torch.maximum, max)
|
|
|
|
def visitStepFunc(self, ctx):
|
|
# step(x, edge) = 1 if x >= edge else 0
|
|
return self._bin_op(
|
|
self.visit(ctx.expr(0)),
|
|
self.visit(ctx.expr(1)),
|
|
lambda x, edge: torch.where(x >= edge, 1.0, 0.0),
|
|
lambda x, edge: 1.0 if x >= edge else 0.0,
|
|
)
|
|
|
|
def visitTopkFunc(self, ctx):
|
|
val = self.visit(ctx.expr(0))
|
|
k = self.visit(ctx.expr(1))
|
|
|
|
if self._is_tensor(k):
|
|
k_val = int(k.flatten()[0].item())
|
|
else:
|
|
k_val = int(k)
|
|
|
|
if self._is_tensor(val):
|
|
size = val.shape[-1]
|
|
k_val = max(1, min(k_val, size))
|
|
score_val = val.abs() if torch.is_complex(val) else val
|
|
|
|
_, indices = torch.topk(score_val, k=k_val, dim=-1)
|
|
indices = indices.contiguous()
|
|
mask = torch.zeros_like(score_val, dtype=torch.bool)
|
|
mask.scatter_(dim=-1, index=indices, value=True)
|
|
|
|
result = torch.where(mask, val, torch.zeros_like(val))
|
|
return result.contiguous()
|
|
|
|
if self._is_list(val):
|
|
k_val = max(0, min(k_val, len(val)))
|
|
try:
|
|
return sorted(val, reverse=True)[:k_val]
|
|
except:
|
|
return val[:k_val]
|
|
|
|
return val
|
|
|
|
def visitBotkFunc(self, ctx):
|
|
val = self.visit(ctx.expr(0))
|
|
k = self.visit(ctx.expr(1))
|
|
|
|
if self._is_tensor(k):
|
|
k_val = int(k.flatten()[0].item())
|
|
else:
|
|
k_val = int(k)
|
|
|
|
if self._is_tensor(val):
|
|
size = val.shape[-1]
|
|
k_val = max(1, min(k_val, size))
|
|
score_val = val.abs() if torch.is_complex(val) else val
|
|
|
|
_, indices = torch.topk(score_val, k=k_val, dim=-1, largest=False)
|
|
indices = indices.contiguous()
|
|
mask = torch.zeros_like(score_val, dtype=torch.bool)
|
|
mask.scatter_(dim=-1, index=indices, value=True)
|
|
|
|
result = torch.where(mask, val, torch.zeros_like(val))
|
|
return result.contiguous()
|
|
|
|
if self._is_list(val):
|
|
k_val = max(0, min(k_val, len(val)))
|
|
try:
|
|
return sorted(val)[:k_val]
|
|
except:
|
|
return val[:k_val]
|
|
|
|
return val
|
|
|
|
def visitPinvFunc(self, ctx):
|
|
"""Permutation inverse: if input[i] = j, output[j] = i."""
|
|
val = self.visit(ctx.expr())
|
|
|
|
if self._is_list(val):
|
|
n = len(val)
|
|
result = [0] * n
|
|
for i, v in enumerate(val):
|
|
idx = int(v)
|
|
if 0 <= idx < n:
|
|
result[idx] = i
|
|
return result
|
|
|
|
if self._is_tensor(val):
|
|
val_list = val.flatten().tolist()
|
|
n = len(val_list)
|
|
result = [0] * n
|
|
for i, v in enumerate(val_list):
|
|
idx = int(v)
|
|
if 0 <= idx < n:
|
|
result[idx] = i
|
|
return torch.tensor(result, device=val.device, dtype=val.dtype).reshape(val.shape)
|
|
|
|
return val
|
|
|
|
# Three-argument functions
|
|
def visitClampFunc(self, ctx):
|
|
val = self.visit(ctx.expr(0))
|
|
min_v = self.visit(ctx.expr(1))
|
|
max_v = self.visit(ctx.expr(2))
|
|
if any(self._is_tensor(x) for x in [val, min_v, max_v]):
|
|
return torch.clamp(self._promote_to_tensor(val), self._promote_to_tensor(min_v), self._promote_to_tensor(max_v))
|
|
if self._is_list(val):
|
|
return [max(min(x, max_v), min_v) for x in val] # TODO: what if min_v or max_v is list
|
|
return max(min(val, max_v), min_v)
|
|
|
|
def visitLerpFunc(self, ctx):
|
|
a = self.visit(ctx.expr(0))
|
|
b = self.visit(ctx.expr(1))
|
|
w = self.visit(ctx.expr(2))
|
|
if any(self._is_tensor(x) for x in [a, b, w]):
|
|
# Lerp: a + w*(b-a)
|
|
return torch.lerp(self._promote_to_tensor(a), self._promote_to_tensor(b), self._promote_to_tensor(w))
|
|
return a + w * (b - a)
|
|
|
|
def visitSmoothstepFunc(self, ctx):
|
|
x = self.visit(ctx.expr(0))
|
|
edge0 = self.visit(ctx.expr(1))
|
|
edge1 = self.visit(ctx.expr(2))
|
|
|
|
# t = clamp((x - edge0) / (edge1 - edge0), 0.0, 1.0)
|
|
# return t * t * (3.0 - 2.0 * t)
|
|
|
|
if any(self._is_tensor(v) for v in [x, edge0, edge1]):
|
|
x_t, e0_t, e1_t = self._promote_to_tensor(x), self._promote_to_tensor(edge0), self._promote_to_tensor(edge1)
|
|
t = torch.clamp((x_t - e0_t) / (e1_t - e0_t), 0.0, 1.0)
|
|
return t * t * (3.0 - 2.0 * t)
|
|
|
|
t = max(0.0, min(1.0, (x - edge0) / (edge1 - edge0)))
|
|
return t * t * (3.0 - 2.0 * t)
|
|
|
|
def visitRangeFunc(self, ctx):
|
|
return list(torch.arange(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), self.visit(ctx.expr(2))))
|
|
|
|
# Helpers for visiting generic exprs
|
|
def visitFunc1Exp(self, ctx):
|
|
return self.visitChildren(ctx)
|
|
|
|
def visitFunc2Exp(self, ctx):
|
|
return self.visitChildren(ctx)
|
|
|
|
def visitFuncNExp(self, ctx):
|
|
return self.visitChildren(ctx)
|
|
|
|
def visitAtomExp(self, ctx):
|
|
return self.visitChildren(ctx)
|
|
|
|
def visitExpr(self, ctx):
|
|
return self.visitChildren(ctx)
|
|
|
|
def visitSMinFunc(self, ctx):
|
|
vals = [self.visit(e) for e in ctx.expr()]
|
|
|
|
if all(not self._is_tensor(x) and not self._is_list(x) for x in vals):
|
|
return min(vals)
|
|
|
|
promoted = [self._promote_to_tensor(x) for x in vals]
|
|
if len(promoted) == 1:
|
|
return torch.min(promoted[0])
|
|
return torch.min(torch.stack(torch.broadcast_tensors(*promoted)))
|
|
|
|
def visitSMaxFunc(self, ctx):
|
|
args = [self.visit(e) for e in ctx.expr()]
|
|
if len(args) == 1:
|
|
# Check if list or scalar
|
|
if not self._is_tensor(args[0]) and not self._is_list(args[0]):
|
|
return args[0]
|
|
if self._is_list(args[0]):
|
|
return max(args[0]) # max of list
|
|
return torch.max(args[0]) # Global max of single tensor
|
|
|
|
# Multiple args
|
|
if all(not self._is_tensor(x) and not self._is_list(x) for x in args):
|
|
return max(args)
|
|
|
|
promoted = [self._promote_to_tensor(x) for x in args]
|
|
if len(promoted) == 1:
|
|
return torch.max(promoted[0])
|
|
return torch.max(torch.stack(torch.broadcast_tensors(*promoted)))
|
|
|
|
def _fold_nd(self, tsr, spatial_dims):
|
|
original_shape = tsr.shape
|
|
added_dims = 0
|
|
target_rank = spatial_dims + 2
|
|
while tsr.ndim < target_rank:
|
|
tsr = tsr.unsqueeze(1)
|
|
added_dims += 1
|
|
folded = False
|
|
if tsr.ndim > target_rank:
|
|
fold_count = tsr.ndim - target_rank
|
|
new_batch = 1
|
|
for i in range(fold_count + 1):
|
|
new_batch *= tsr.shape[i]
|
|
tsr = tsr.reshape(new_batch, *tsr.shape[fold_count + 1 :])
|
|
folded = True
|
|
else:
|
|
folded = added_dims > 0
|
|
return tsr, original_shape, added_dims, folded
|
|
|
|
def _unfold_nd(self, tsr, original_shape, added_dims, folded):
|
|
spatial_dims = tsr.ndim - 2
|
|
if folded and added_dims == 0:
|
|
target_fold_rank = len(original_shape) - (spatial_dims + 1)
|
|
fold_dims = original_shape[:target_fold_rank]
|
|
tsr = tsr.reshape(*fold_dims, *tsr.shape[1:])
|
|
for _ in range(added_dims):
|
|
if tsr.ndim > len(original_shape) and tsr.shape[1] == 1:
|
|
tsr = tsr.squeeze(1)
|
|
return tsr
|
|
|
|
def visitPermuteFunc(self, ctx):
|
|
tsr = self._promote_to_tensor(self.visit(ctx.expr(0)))
|
|
dims = self.visit(ctx.expr(1))
|
|
if isinstance(dims, torch.Tensor):
|
|
dims = dims.flatten().long().tolist()
|
|
return tsr.permute(*dims)
|
|
|
|
def visitReshapeFunc(self, ctx):
|
|
tsr = self._promote_to_tensor(self.visit(ctx.expr(0)))
|
|
new_shape = self.visit(ctx.expr(1))
|
|
if isinstance(new_shape, torch.Tensor):
|
|
new_shape = new_shape.flatten().long().tolist()
|
|
return tsr.reshape(*new_shape)
|
|
|
|
def visitPrintShapeFunc(self, ctx):
|
|
tsr = self.visit(ctx.expr())
|
|
if hasattr(tsr, "shape"):
|
|
print(tsr.shape)
|
|
else:
|
|
print(f"Scalar/List: {tsr}")
|
|
return tsr
|
|
|
|
def visitSfftFunc(self, ctx):
|
|
old_vars = self.variables
|
|
self.variables = self.spatial_variables.copy()
|
|
try:
|
|
val = self._promote_to_tensor(self.visit(ctx.expr()))
|
|
dims = tuple(range(val.ndim))
|
|
return torch.fft.fftn(val, dim=dims)
|
|
finally:
|
|
self.variables = old_vars
|
|
|
|
def visitSifftFunc(self, ctx):
|
|
old_vars = self.variables
|
|
self.variables = self.variables.copy()
|
|
device = self.device
|
|
|
|
shape_to_use = self.shape if self.shape else (1, 1, 1, 1)
|
|
|
|
ndim = len(shape_to_use)
|
|
dim_names = ["x", "y", "z", "w", "v", "u"]
|
|
|
|
k_components = []
|
|
for i in range(ndim):
|
|
dim_idx = ndim - 1 - i
|
|
size_d = shape_to_use[dim_idx]
|
|
values = torch.arange(size_d, dtype=torch.float32, device=device)
|
|
view_shape = [1] * ndim
|
|
view_shape[dim_idx] = size_d
|
|
values = values.view(*view_shape).expand(*shape_to_use)
|
|
|
|
if i < len(dim_names):
|
|
var_name = f"K{dim_names[i]}"
|
|
self.variables[var_name] = values
|
|
self.variables[f"F{dim_names[i]}"] = float(size_d)
|
|
|
|
self.variables[f"K_dim{dim_idx}"] = values
|
|
self.variables[f"F_dim{dim_idx}"] = float(size_d)
|
|
k_components.append(values)
|
|
|
|
k_sq_sum = torch.zeros(shape_to_use, device=device)
|
|
for k_val in k_components:
|
|
k_sq_sum = k_sq_sum + k_val**2
|
|
|
|
self.variables["K"] = torch.sqrt(k_sq_sum)
|
|
self.variables["frequency"] = self.variables["K"]
|
|
|
|
if "Kx" in self.variables:
|
|
self.variables["frequency_count"] = self.variables.get("Fx", 1.0)
|
|
|
|
self.variables = self.variables | generate_dim_variables(k_sq_sum)
|
|
|
|
try:
|
|
val = self._promote_to_tensor(self.visit(ctx.expr()))
|
|
dims = tuple(range(val.ndim))
|
|
return torch.fft.ifftn(val, dim=dims).real
|
|
finally:
|
|
self.variables = old_vars
|
|
|
|
def visitSwapFunc(self, ctx):
|
|
tsr = self._promote_to_tensor(self.visit(ctx.expr(0)))
|
|
dim_t = self.visit(ctx.expr(1))
|
|
idx1_t = self.visit(ctx.expr(2))
|
|
idx2_t = self.visit(ctx.expr(3))
|
|
|
|
dim = int(dim_t.flatten()[0].item()) if isinstance(dim_t, torch.Tensor) else int(dim_t)
|
|
i = int(idx1_t.flatten()[0].item()) if isinstance(idx1_t, torch.Tensor) else int(idx1_t)
|
|
j = int(idx2_t.flatten()[0].item()) if isinstance(idx2_t, torch.Tensor) else int(idx2_t)
|
|
|
|
while dim < 0:
|
|
dim += tsr.ndim
|
|
while i < 0:
|
|
i += tsr.shape[dim]
|
|
while j < 0:
|
|
j += tsr.shape[dim]
|
|
|
|
indices = torch.arange(tsr.shape[dim], device=tsr.device)
|
|
val_i = indices[i].clone()
|
|
indices[i] = indices[j]
|
|
indices[j] = val_i
|
|
|
|
return torch.index_select(tsr, dim, indices)
|
|
|
|
def _normalize_coord(self, coord, size):
|
|
if size > 1:
|
|
return (coord / (size - 1)) * 2.0 - 1.0
|
|
return torch.zeros_like(coord)
|
|
|
|
def visitSumFunc(self, ctx):
|
|
return self._reduction_op(self.visit(ctx.expr()), torch.sum, sum)
|
|
|
|
def visitMeanFunc(self, ctx):
|
|
return self._reduction_op(
|
|
self.visit(ctx.expr()), lambda x: torch.mean(x.float()), lambda x: sum(x) / len(x) if x else 0.0
|
|
)
|
|
|
|
def visitStdFunc(self, ctx):
|
|
def list_std(val):
|
|
if len(val) < 2:
|
|
return 0.0
|
|
mean = sum(val) / len(val)
|
|
variance = sum((x - mean) ** 2 for x in val) / (len(val) - 1)
|
|
return math.sqrt(variance)
|
|
|
|
return self._reduction_op(self.visit(ctx.expr()), lambda x: torch.std(x.float()), list_std)
|
|
|
|
def visitVarFunc(self, ctx):
|
|
def list_var(val):
|
|
if len(val) < 2:
|
|
return 0.0
|
|
mean = sum(val) / len(val)
|
|
return sum((x - mean) ** 2 for x in val) / (len(val) - 1)
|
|
|
|
return self._reduction_op(self.visit(ctx.expr()), lambda x: torch.var(x.float()), list_var)
|
|
|
|
def _manual_quantile(self, val, q):
|
|
"""Fallback implementation using sort for when torch.quantile fails on large tensors."""
|
|
val_flat = val.flatten().float()
|
|
sorted_val, _ = torch.sort(val_flat)
|
|
n = len(sorted_val)
|
|
if n == 0:
|
|
return torch.zeros_like(q) if self._is_tensor(q) else 0.0
|
|
|
|
# indices = q * (n - 1)
|
|
indices = q * (n - 1)
|
|
low = torch.floor(indices).long()
|
|
high = torch.ceil(indices).long()
|
|
frac = (indices - low).float()
|
|
|
|
# Ensure bounds
|
|
low = torch.clamp(low, 0, n - 1)
|
|
high = torch.clamp(high, 0, n - 1)
|
|
|
|
res = sorted_val[low] + (sorted_val[high] - sorted_val[low]) * frac
|
|
return res
|
|
|
|
def _quartile_helper(self, val, q):
|
|
|
|
if self._is_tensor(q) and not self._is_tensor(val):
|
|
val = self._promote_to_tensor(val)
|
|
|
|
if self._is_tensor(val):
|
|
if not self._is_tensor(q):
|
|
q = torch.tensor(q, device=self.device).float()
|
|
|
|
try:
|
|
if q.ndim > 1:
|
|
q_flat = q.flatten()
|
|
res = torch.quantile(val.float(), q_flat)
|
|
return res.reshape(q.shape)
|
|
|
|
return torch.quantile(val.float(), q)
|
|
except RuntimeError as e:
|
|
# Fallback for "input tensor is too large" or other quantile-specific issues
|
|
if "quantile" in str(e).lower() or "too large" in str(e).lower():
|
|
if q.ndim > 1:
|
|
q_flat = q.flatten()
|
|
res = self._manual_quantile(val, q_flat)
|
|
return res.reshape(q.shape)
|
|
return self._manual_quantile(val, q)
|
|
raise e
|
|
|
|
if self._is_list(val):
|
|
if not val:
|
|
return 0.0
|
|
sorted_data = sorted(val)
|
|
n = len(val)
|
|
pos = (n - 1) * q
|
|
whole = int(pos)
|
|
frac = pos - whole
|
|
if whole + 1 < n:
|
|
return sorted_data[whole] + (sorted_data[whole + 1] - sorted_data[whole]) * frac
|
|
else:
|
|
return sorted_data[whole]
|
|
return val
|
|
|
|
def visitQuartileFunc(self, ctx):
|
|
val = self.visit(ctx.expr(0))
|
|
k = self.visit(ctx.expr(1))
|
|
|
|
if self._is_tensor(k):
|
|
# q = k * 0.25. Ensure k is treated as int-like (1,2,3)?
|
|
# Old code did int().
|
|
return self._quartile_helper(val, (k.int().float() * 0.25))
|
|
if self._is_list(k):
|
|
return [self._quartile_helper(val, int(x) * 0.25) for x in k]
|
|
|
|
k_val = int(k)
|
|
q = min(1.0, max(0.0, k_val * 0.25))
|
|
return self._quartile_helper(val, q)
|
|
|
|
def visitPercentileFunc(self, ctx):
|
|
val = self.visit(ctx.expr(0))
|
|
p_raw = self.visit(ctx.expr(1))
|
|
|
|
if self._is_tensor(p_raw):
|
|
# p is 0-100. q = p / 100
|
|
return self._quartile_helper(val, p_raw.float() / 100.0)
|
|
if self._is_list(p_raw):
|
|
return [self._quartile_helper(val, float(x) / 100.0) for x in p_raw]
|
|
|
|
p = float(p_raw)
|
|
q = max(0.0, min(1.0, p / 100.0))
|
|
return self._quartile_helper(val, q)
|
|
|
|
def visitQuantileFunc(self, ctx):
|
|
val = self.visit(ctx.expr(0))
|
|
q_raw = self.visit(ctx.expr(1))
|
|
|
|
if self._is_tensor(q_raw):
|
|
return self._quartile_helper(val, q_raw.float())
|
|
if self._is_list(q_raw):
|
|
return [self._quartile_helper(val, float(x)) for x in q_raw]
|
|
|
|
q = float(q_raw)
|
|
q = max(0.0, min(1.0, q))
|
|
return self._quartile_helper(val, q)
|
|
|
|
def visitDotFunc(self, ctx):
|
|
a = self._promote_to_tensor(self.visit(ctx.expr(0)))
|
|
b = self._promote_to_tensor(self.visit(ctx.expr(1)))
|
|
return torch.dot(a.flatten(), b.flatten())
|
|
|
|
def visitMomentFunc(self,ctx):
|
|
x = self._promote_to_tensor(self.visit(ctx.expr(0)))
|
|
a = self.visit(ctx.expr(1))
|
|
k = self.visit(ctx.expr(2))
|
|
|
|
return torch.sum(self._bin_op(self._bin_op(x,a,torch.sub,lambda x, a: x - a),k,torch.pow,pow)).item()/x.numel()
|
|
|
|
def visitMapFunc(self, ctx):
|
|
tensor = self._promote_to_tensor(self.visit(ctx.expr(0)))
|
|
coords = [self._promote_to_tensor(self.visit(ctx.expr(i))) for i in range(1, len(ctx.expr()))]
|
|
num_coords = len(coords)
|
|
|
|
if num_coords == 0:
|
|
return tensor
|
|
if num_coords > 3:
|
|
raise ValueError("map() supports max 3 mapping functions.")
|
|
|
|
spatial_in_shape = tensor.shape[-num_coords:]
|
|
leading_shape = tensor.shape[:-num_coords]
|
|
|
|
batch_size = 1
|
|
for s in leading_shape:
|
|
batch_size *= s
|
|
|
|
input_view = tensor.reshape(batch_size, 1, *spatial_in_shape)
|
|
norm_coords_list = []
|
|
for i in range(num_coords):
|
|
dim_size = spatial_in_shape[i]
|
|
norm = self._normalize_coord(coords[i], dim_size)
|
|
norm_coords_list.append(norm)
|
|
|
|
grid = torch.stack(norm_coords_list[::-1], dim=-1)
|
|
grid_spatial_shape = grid.shape[:-1]
|
|
|
|
try:
|
|
grid_view = grid.reshape(batch_size, *grid_spatial_shape[-(num_coords):], num_coords)
|
|
except RuntimeError:
|
|
grid_view = grid.expand(batch_size, *([-1] * len(grid_spatial_shape)), -1)
|
|
grid_view = grid_view.reshape(batch_size, *grid_view.shape[-(num_coords + 1) : -1], num_coords)
|
|
|
|
if num_coords == 1:
|
|
input_final = input_view.reshape(batch_size, 1, 1, -1)
|
|
gv = grid_view
|
|
while gv.ndim < 3:
|
|
gv = gv.unsqueeze(1)
|
|
y_zeros = torch.zeros_like(gv[..., :1])
|
|
grid_final = torch.cat([gv, y_zeros], dim=-1).unsqueeze(1) # [B, 1, 1, 2]
|
|
output = F.grid_sample(input_final, grid_final, align_corners=True)
|
|
elif num_coords == 2:
|
|
grid_final = grid_view.reshape(batch_size, *grid_view.shape[-3:-1], 2)
|
|
output = F.grid_sample(input_view, grid_final, align_corners=True)
|
|
else:
|
|
grid_final = grid_view.reshape(batch_size, *grid_view.shape[-4:-1], 3)
|
|
output = F.grid_sample(input_view, grid_final, align_corners=True)
|
|
actual_spatial = grid_view.shape[1:-1]
|
|
final_shape = list(leading_shape) + list(actual_spatial)
|
|
return output.reshape(final_shape)
|
|
|
|
def _parse_conv_args(self, ctx):
|
|
input_raw = self.visit(ctx.expr(0))
|
|
tensor = self._promote_to_tensor(input_raw)
|
|
num_args = len(ctx.expr())
|
|
if num_args < 3:
|
|
raise ValueError("conv() requires at least 3 arguments")
|
|
|
|
kernel_arg_idx = num_args - 1
|
|
|
|
spatial_dims_count = num_args - 2
|
|
if spatial_dims_count not in [1, 2, 3]:
|
|
raise ValueError(f"conv() supports 1D, 2D, or 3D. Found {spatial_dims_count}")
|
|
|
|
kernel_sizes = []
|
|
for i in range(1, 1 + spatial_dims_count):
|
|
val = self.visit(ctx.expr(i))
|
|
if isinstance(val, torch.Tensor):
|
|
val = int(val.flatten()[0].item())
|
|
kernel_sizes.append(val)
|
|
|
|
# Prepare context for kernel
|
|
coords = [torch.arange(s, device=self.device).float() - (s // 2) for s in kernel_sizes]
|
|
grid = torch.meshgrid(*coords, indexing="ij")
|
|
|
|
dim_names = ["x", "y", "z"]
|
|
old_vars = self.variables
|
|
self.variables = self.variables.copy()
|
|
for i in range(spatial_dims_count):
|
|
self.variables[f"k{dim_names[i].lower()}"] = grid[i]
|
|
self.variables[f"k{dim_names[i].upper()}"] = grid[i]
|
|
|
|
kernel_var_names = ["kW", "kH", "kD"]
|
|
kernel_var_full = ["kernel_width", "kernel_height", "kernel_depth"]
|
|
for i in range(spatial_dims_count):
|
|
size_val = float(kernel_sizes[i])
|
|
self.variables[kernel_var_names[i]] = size_val
|
|
self.variables[kernel_var_full[i]] = size_val
|
|
|
|
original_shape = self.shape
|
|
self.shape = tuple(kernel_sizes)
|
|
|
|
try:
|
|
kernel_val = self._promote_to_tensor(self.visit(ctx.expr(kernel_arg_idx)))
|
|
finally:
|
|
self.shape = original_shape
|
|
self.variables = old_vars
|
|
|
|
return tensor, kernel_val, kernel_sizes, spatial_dims_count
|
|
|
|
def _apply_conv_internal(self, conv_input, kernel_val, kernel_sizes, spatial_dims_count):
|
|
"""
|
|
Executes the convolution with asymmetric padding support for even kernels.
|
|
Input: [Batch, Channel, Spatial...]
|
|
Kernel: [Spatial...] (to be promoted/repeated)
|
|
kernel_sizes: [W, H, D]
|
|
"""
|
|
in_channels = conv_input.size(1)
|
|
|
|
pads = []
|
|
for k in kernel_sizes:
|
|
p_total = k - 1
|
|
p_low = k // 2
|
|
p_high = p_total - p_low
|
|
pads.extend([p_low, p_high])
|
|
|
|
padded_input = torch.nn.functional.pad(conv_input, tuple(pads), mode="constant", value=0)
|
|
|
|
actual_kernel_sizes = kernel_sizes[::-1]
|
|
|
|
if kernel_val.numel() == 1:
|
|
kernel_val = kernel_val.expand(tuple(actual_kernel_sizes))
|
|
elif kernel_val.ndim != spatial_dims_count:
|
|
kernel_val = kernel_val.reshape(tuple(actual_kernel_sizes))
|
|
|
|
final_kernel = kernel_val.unsqueeze(0).unsqueeze(0)
|
|
final_kernel = final_kernel.to(conv_input.dtype)
|
|
final_kernel = final_kernel.repeat(in_channels, 1, *([1] * spatial_dims_count))
|
|
|
|
conv_fn = F.conv1d if spatial_dims_count == 1 else (F.conv2d if spatial_dims_count == 2 else F.conv3d)
|
|
|
|
result = conv_fn(padded_input, final_kernel, padding=0, groups=in_channels)
|
|
return result
|
|
|
|
def visitEzConvFunc(self, ctx):
|
|
tensor, kernel_val, kernel_sizes, spatial_dims_count = self._parse_conv_args(ctx)
|
|
input_ndim = tensor.ndim
|
|
|
|
is_channels_first = False
|
|
if tensor.ndim == spatial_dims_count + 2:
|
|
c_front = tensor.shape[1]
|
|
c_back = tensor.shape[-1]
|
|
if c_front <= 4 and c_front < c_back:
|
|
is_channels_first = True
|
|
elif c_back <= 4 and c_back < c_front:
|
|
is_channels_first = False
|
|
else:
|
|
is_channels_first = c_back > 4
|
|
|
|
if is_channels_first:
|
|
if spatial_dims_count == 1 and tensor.ndim == 3:
|
|
tensor = tensor.permute(0, 2, 1)
|
|
elif spatial_dims_count == 2 and tensor.ndim == 4:
|
|
tensor = tensor.permute(0, 2, 3, 1)
|
|
elif spatial_dims_count == 3:
|
|
if tensor.ndim == 4:
|
|
tensor = tensor.unsqueeze(-1)
|
|
elif tensor.ndim == 5:
|
|
tensor = tensor.permute(0, 2, 3, 4, 1)
|
|
|
|
if tensor.ndim == spatial_dims_count + 1:
|
|
tensor = tensor.unsqueeze(-1) # Add C=1
|
|
elif tensor.ndim == spatial_dims_count:
|
|
tensor = tensor.unsqueeze(0).unsqueeze(-1) # Add B=1, C=1
|
|
|
|
in_channels = tensor.size(-1)
|
|
batch_end_idx = tensor.ndim - 1 - spatial_dims_count
|
|
batch_shape = tensor.shape[:batch_end_idx]
|
|
spatial_shape = tensor.shape[batch_end_idx:-1]
|
|
channels_shape = (tensor.shape[-1],)
|
|
|
|
total_batch = 1
|
|
for s in batch_shape:
|
|
total_batch *= s
|
|
flat_input = tensor.reshape(total_batch, *spatial_shape, in_channels)
|
|
|
|
permute_order = [0, spatial_dims_count + 1] + list(range(1, spatial_dims_count + 1))
|
|
conv_input = flat_input.permute(*permute_order)
|
|
|
|
result = self._apply_conv_internal(conv_input, kernel_val, kernel_sizes, spatial_dims_count)
|
|
|
|
result_permute = [0] + list(range(2, 2 + spatial_dims_count)) + [1]
|
|
out_flat = result.permute(*result_permute)
|
|
|
|
final_shape = batch_shape + spatial_shape + channels_shape
|
|
out = out_flat.reshape(final_shape)
|
|
|
|
if is_channels_first:
|
|
if spatial_dims_count == 1 and out.ndim == 3:
|
|
out = out.permute(0, 2, 1)
|
|
elif spatial_dims_count == 2 and out.ndim == 4:
|
|
out = out.permute(0, 3, 1, 2)
|
|
elif spatial_dims_count == 3:
|
|
if out.ndim == 5 and out.shape[-1] == 1 and input_ndim == 4:
|
|
out = out.squeeze(-1)
|
|
elif out.ndim == 5:
|
|
out = out.permute(0, 4, 1, 2, 3)
|
|
|
|
return out
|
|
|
|
def visitConvFunc(self, ctx):
|
|
tensor, kernel_val, kernel_sizes, spatial_dims_count = self._parse_conv_args(ctx)
|
|
|
|
# Expect (Batch..., Channel, Spatial...)
|
|
# spatial_dims_count = 1, 2, or 3
|
|
# Must have at least Channel + Spatial dims
|
|
min_dims = spatial_dims_count + 1
|
|
if tensor.ndim < min_dims:
|
|
raise ValueError(f"convolution() input requires at least Channels + Spatial dimensions. Got shape {tensor.shape} for {spatial_dims_count}D conv.")
|
|
|
|
spatial_shape = tensor.shape[-spatial_dims_count:]
|
|
in_channels = tensor.shape[-(spatial_dims_count + 1)]
|
|
batch_shape = tensor.shape[:-(spatial_dims_count + 1)]
|
|
|
|
total_batch = 1
|
|
for s in batch_shape:
|
|
total_batch *= s
|
|
|
|
conv_input = tensor.reshape(total_batch, in_channels, *spatial_shape)
|
|
|
|
result = self._apply_conv_internal(conv_input, kernel_val, kernel_sizes, spatial_dims_count)
|
|
|
|
return result.reshape(*batch_shape, in_channels, *spatial_shape)
|