ai: slicing
This commit is contained in:
mcDandy
@@ -80,9 +80,20 @@ unaryExpr:
|
||||
| indexExpr # ToIndex;
|
||||
|
||||
indexExpr:
|
||||
indexExpr LBRACKET expr (COMMA expr)* RBRACKET # IndexExp
|
||||
indexExpr LBRACKET sliceElement (COMMA sliceElement)* RBRACKET # IndexExp
|
||||
| atom # ToAtom;
|
||||
|
||||
sliceElement:
|
||||
expr COLON expr COLON expr # SliceWithStep
|
||||
| expr COLON expr # Slice
|
||||
| COLON expr COLON expr # SliceColonStartStep
|
||||
| COLON expr # SliceColonStart
|
||||
| expr COLON COLON expr # SliceExprColonStep
|
||||
| expr COLON # SliceExpr
|
||||
| COLON COLON expr # SliceColonStep
|
||||
| COLON # SliceOnly
|
||||
| expr # IndexOnly;
|
||||
|
||||
// Atoms: function calls, variable, number, constant, or parenthesized expression
|
||||
atom:
|
||||
func0 # Func0Exp
|
||||
|
||||
@@ -455,37 +455,92 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
def visitToIndex(self, ctx):
|
||||
return (yield ctx.indexExpr())
|
||||
|
||||
def visitIndexExp(self, ctx):
|
||||
val = (yield ctx.indexExpr())
|
||||
raw_index_nodes = ctx.expr()
|
||||
|
||||
indices = []
|
||||
for node in raw_index_nodes:
|
||||
idx_val = (yield node)
|
||||
def _process_slice_element(self, slice_elem_ctx):
|
||||
"""Process a sliceElement context and return either an index (int) or a slice object."""
|
||||
# Get the context class name to determine which type of element this is
|
||||
ctx_type = type(slice_elem_ctx).__name__
|
||||
|
||||
if ctx_type == "IndexOnlyContext":
|
||||
# Regular index: expr
|
||||
idx_val = yield slice_elem_ctx.expr(0)
|
||||
if self._is_tensor(idx_val):
|
||||
if idx_val.numel() == 1:
|
||||
indices.append(int(idx_val.flatten()[0].item()))
|
||||
return int(idx_val.flatten()[0].item())
|
||||
else:
|
||||
# Fancy indexing with tensor
|
||||
indices.append(idx_val.long())
|
||||
return idx_val.long()
|
||||
elif self._is_list(idx_val):
|
||||
# Fancy indexing with list - convert to tensor
|
||||
indices.append(torch.tensor(idx_val, dtype=torch.long, device=self.device))
|
||||
return torch.tensor(idx_val, dtype=torch.long, device=self.device)
|
||||
else:
|
||||
indices.append(int(idx_val))
|
||||
return int(idx_val)
|
||||
|
||||
# Use standard PyTorch/list indexing
|
||||
# Handle slice expressions - all cases with colons
|
||||
exprs = slice_elem_ctx.expr() if hasattr(slice_elem_ctx, 'expr') else []
|
||||
|
||||
start = None
|
||||
stop = None
|
||||
step = None
|
||||
|
||||
if ctx_type == "SliceContext":
|
||||
# expr COLON expr -> [start:stop]
|
||||
start = yield exprs[0]
|
||||
stop = yield exprs[1]
|
||||
elif ctx_type == "SliceWithStepContext":
|
||||
# expr COLON expr COLON expr -> [start:stop:step]
|
||||
start = yield exprs[0]
|
||||
stop = yield exprs[1]
|
||||
step = yield exprs[2]
|
||||
elif ctx_type == "SliceExprContext":
|
||||
# expr COLON -> [start:]
|
||||
start = yield exprs[0]
|
||||
elif ctx_type == "SliceColonStartContext":
|
||||
# COLON expr -> [:stop]
|
||||
stop = yield exprs[0]
|
||||
elif ctx_type == "SliceColonStartStepContext":
|
||||
# COLON expr COLON expr -> [:stop:step]
|
||||
stop = yield exprs[0]
|
||||
step = yield exprs[1]
|
||||
elif ctx_type == "SliceColonStepContext":
|
||||
# COLON COLON expr -> [::step]
|
||||
step = yield exprs[0]
|
||||
elif ctx_type == "SliceExprColonStepContext":
|
||||
# expr COLON COLON expr -> [start::step]
|
||||
start = yield exprs[0]
|
||||
step = yield exprs[1]
|
||||
elif ctx_type == "SliceOnlyContext":
|
||||
# COLON -> [:]
|
||||
pass
|
||||
|
||||
# Normalize start, stop, step to int or None
|
||||
start_val = None if start is None else int(start) if not isinstance(start, torch.Tensor) else int(start.item())
|
||||
stop_val = None if stop is None else int(stop) if not isinstance(stop, torch.Tensor) else int(stop.item())
|
||||
step_val = None if step is None else int(step) if not isinstance(step, torch.Tensor) else int(step.item())
|
||||
|
||||
return slice(start_val, stop_val, step_val)
|
||||
|
||||
def visitIndexExp(self, ctx):
|
||||
val = (yield ctx.indexExpr())
|
||||
slice_elem_nodes = ctx.sliceElement()
|
||||
|
||||
indices = []
|
||||
for node in slice_elem_nodes:
|
||||
elem = yield self._process_slice_element(node)
|
||||
indices.append(elem)
|
||||
|
||||
# Use standard PyTorch/list indexing and slicing
|
||||
if self._is_tensor(val):
|
||||
if len(indices) > val.ndim:
|
||||
raise ValueError(f"{ctx.start.line}:{ctx.start.column}: Expacted up to {val.ndim} dimensions but got {indices}.")
|
||||
raise ValueError(f"{ctx.start.line}:{ctx.start.column}: Expected up to {val.ndim} dimensions but got {len(indices)}.")
|
||||
|
||||
normalized_indices = []
|
||||
# Build index tuple, handling both regular indices and slices
|
||||
idx_tuple = []
|
||||
for dim, idx in enumerate(indices):
|
||||
normalized_indices.append(self._normalize_index_value(idx, val.shape[dim], ctx, f"dimension {dim}"))
|
||||
if isinstance(idx, slice):
|
||||
idx_tuple.append(idx)
|
||||
else:
|
||||
# Normalize scalar index
|
||||
idx_tuple.append(self._normalize_index_value(idx, val.shape[dim], ctx, f"dimension {dim}"))
|
||||
|
||||
idx_tuple = tuple(normalized_indices)
|
||||
result = val[idx_tuple]
|
||||
result = val[tuple(idx_tuple)]
|
||||
if self._is_tensor(result):
|
||||
if getattr(result, "is_nested", False):
|
||||
return result
|
||||
@@ -496,27 +551,73 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
elif isinstance(val, str):
|
||||
current = val
|
||||
for idx in indices:
|
||||
if isinstance(idx, torch.Tensor):
|
||||
if isinstance(idx, slice):
|
||||
current = current[idx]
|
||||
elif isinstance(idx, torch.Tensor):
|
||||
if idx.numel() != 1:
|
||||
raise ValueError(f"{ctx.start.line}:{ctx.start.column}: Too many indecies for string with 1 dimension. Got {idx.numel()}")
|
||||
raise ValueError(f"{ctx.start.line}:{ctx.start.column}: Too many indices for string with 1 dimension. Got {idx.numel()}")
|
||||
idx = int(idx.flatten()[0].item())
|
||||
idx = self._normalize_index_value(idx, len(current), ctx, "string")
|
||||
current = current[idx]
|
||||
idx = self._normalize_index_value(idx, len(current), ctx, "string")
|
||||
current = current[idx]
|
||||
else:
|
||||
idx = self._normalize_index_value(idx, len(current), ctx, "string")
|
||||
current = current[idx]
|
||||
return current
|
||||
elif self._is_list(val):
|
||||
# Navigate through nested lists
|
||||
# Navigate through nested lists, supporting both indexing and slicing
|
||||
current = val
|
||||
for idx in indices:
|
||||
if isinstance(idx, torch.Tensor):
|
||||
if isinstance(idx, slice):
|
||||
current = current[idx]
|
||||
elif isinstance(idx, torch.Tensor):
|
||||
if idx.numel() != 1:
|
||||
raise ValueError(f"{ctx.start.line}:{ctx.start.column}: List index must be a scalar. Got tensor with length of {idx.numel()}")
|
||||
idx = int(idx.item())
|
||||
idx = self._normalize_index_value(idx, len(current), ctx, "list")
|
||||
current = current[idx]
|
||||
idx = self._normalize_index_value(idx, len(current), ctx, "list")
|
||||
current = current[idx]
|
||||
else:
|
||||
idx = self._normalize_index_value(idx, len(current), ctx, "list")
|
||||
current = current[idx]
|
||||
return current
|
||||
error_prefix = f"{ctx.start.line}:{ctx.start.column}:"
|
||||
raise ValueError(f"{error_prefix} Indexing only supported on tensors, lists, and strings (found {type(val).__name__})")
|
||||
|
||||
def visitIndexOnly(self, ctx):
|
||||
"""Visit slice element that is just an index expression."""
|
||||
return (yield ctx.expr(0))
|
||||
|
||||
def visitSlice(self, ctx):
|
||||
"""Visit slice element with start and stop: expr COLON expr"""
|
||||
return None # Handled by _process_slice_element
|
||||
|
||||
def visitSliceWithStep(self, ctx):
|
||||
"""Visit slice element with start, stop, and step: expr COLON expr COLON expr"""
|
||||
return None # Handled by _process_slice_element
|
||||
|
||||
def visitSliceColonStart(self, ctx):
|
||||
"""Visit slice element with stop only: COLON expr"""
|
||||
return None # Handled by _process_slice_element
|
||||
|
||||
def visitSliceColonStartStep(self, ctx):
|
||||
"""Visit slice element with stop and step: COLON expr COLON expr"""
|
||||
return None # Handled by _process_slice_element
|
||||
|
||||
def visitSliceExpr(self, ctx):
|
||||
"""Visit slice element with start only: expr COLON"""
|
||||
return None # Handled by _process_slice_element
|
||||
|
||||
def visitSliceColonStep(self, ctx):
|
||||
"""Visit slice element with step only: COLON COLON expr"""
|
||||
return None # Handled by _process_slice_element
|
||||
|
||||
def visitSliceExprColonStep(self, ctx):
|
||||
"""Visit slice element with start and step: expr COLON COLON expr"""
|
||||
return None # Handled by _process_slice_element
|
||||
|
||||
def visitSliceOnly(self, ctx):
|
||||
"""Visit slice element with no bounds: COLON"""
|
||||
return None # Handled by _process_slice_element
|
||||
|
||||
def visitToAtom(self, ctx):
|
||||
return (yield ctx.atom())
|
||||
|
||||
@@ -1632,14 +1733,8 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
return float(torch.sum(self._bin_op(self._bin_op(x,a,torch.sub,lambda x, a: x - a,ctx),k,torch.pow,pow,ctx)).item())/x.numel()
|
||||
|
||||
def visitSortFunc(self, ctx):
|
||||
val = self._promote_to_tensor((yield ctx.expr(0)))
|
||||
desc = False
|
||||
dim = -1
|
||||
if len(ctx.expr()) > 1:
|
||||
desc = bool((yield ctx.expr(1)))
|
||||
if len(ctx.expr()) > 2:
|
||||
dim = int((yield ctx.expr(2)))
|
||||
sorted_val, _ = torch.sort(val, descending=desc, dim=dim)
|
||||
val = self._promote_to_tensor((yield ctx.expr()))
|
||||
sorted_val, _ = torch.sort(val)
|
||||
return sorted_val
|
||||
|
||||
def visitCossimFunc(self, ctx):
|
||||
@@ -1653,9 +1748,6 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
except RuntimeError as e:
|
||||
error_msg = f"{ctx.start.line}:{ctx.start.column}: cossim({a.shape}, {b.shape}): Incompatible shapes for cosine similarity - {str(e)}"
|
||||
raise ValueError(error_msg)
|
||||
except ValueError as e:
|
||||
error_msg = f"{ctx.start.line}:{ctx.start.column}: cossim({a.shape}, {b.shape}): {str(e)}"
|
||||
raise ValueError(error_msg)
|
||||
|
||||
def visitRifeFunc(self, ctx):
|
||||
img1 = self._promote_to_tensor((yield ctx.expr(0)))
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -422,6 +422,87 @@ class MathExprListener(ParseTreeListener):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceWithStep.
|
||||
def enterSliceWithStep(self, ctx:MathExprParser.SliceWithStepContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceWithStep.
|
||||
def exitSliceWithStep(self, ctx:MathExprParser.SliceWithStepContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#Slice.
|
||||
def enterSlice(self, ctx:MathExprParser.SliceContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#Slice.
|
||||
def exitSlice(self, ctx:MathExprParser.SliceContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceColonStartStep.
|
||||
def enterSliceColonStartStep(self, ctx:MathExprParser.SliceColonStartStepContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceColonStartStep.
|
||||
def exitSliceColonStartStep(self, ctx:MathExprParser.SliceColonStartStepContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceColonStart.
|
||||
def enterSliceColonStart(self, ctx:MathExprParser.SliceColonStartContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceColonStart.
|
||||
def exitSliceColonStart(self, ctx:MathExprParser.SliceColonStartContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceExprColonStep.
|
||||
def enterSliceExprColonStep(self, ctx:MathExprParser.SliceExprColonStepContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceExprColonStep.
|
||||
def exitSliceExprColonStep(self, ctx:MathExprParser.SliceExprColonStepContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceExpr.
|
||||
def enterSliceExpr(self, ctx:MathExprParser.SliceExprContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceExpr.
|
||||
def exitSliceExpr(self, ctx:MathExprParser.SliceExprContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceColonStep.
|
||||
def enterSliceColonStep(self, ctx:MathExprParser.SliceColonStepContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceColonStep.
|
||||
def exitSliceColonStep(self, ctx:MathExprParser.SliceColonStepContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceOnly.
|
||||
def enterSliceOnly(self, ctx:MathExprParser.SliceOnlyContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceOnly.
|
||||
def exitSliceOnly(self, ctx:MathExprParser.SliceOnlyContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#IndexOnly.
|
||||
def enterIndexOnly(self, ctx:MathExprParser.IndexOnlyContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#IndexOnly.
|
||||
def exitIndexOnly(self, ctx:MathExprParser.IndexOnlyContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#Func0Exp.
|
||||
def enterFunc0Exp(self, ctx:MathExprParser.Func0ExpContext):
|
||||
pass
|
||||
|
||||
+3173
-2765
File diff suppressed because it is too large
Load Diff
@@ -239,6 +239,51 @@ class MathExprVisitor(ParseTreeVisitor):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceWithStep.
|
||||
def visitSliceWithStep(self, ctx:MathExprParser.SliceWithStepContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#Slice.
|
||||
def visitSlice(self, ctx:MathExprParser.SliceContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceColonStartStep.
|
||||
def visitSliceColonStartStep(self, ctx:MathExprParser.SliceColonStartStepContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceColonStart.
|
||||
def visitSliceColonStart(self, ctx:MathExprParser.SliceColonStartContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceExprColonStep.
|
||||
def visitSliceExprColonStep(self, ctx:MathExprParser.SliceExprColonStepContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceExpr.
|
||||
def visitSliceExpr(self, ctx:MathExprParser.SliceExprContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceColonStep.
|
||||
def visitSliceColonStep(self, ctx:MathExprParser.SliceColonStepContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceOnly.
|
||||
def visitSliceOnly(self, ctx:MathExprParser.SliceOnlyContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#IndexOnly.
|
||||
def visitIndexOnly(self, ctx:MathExprParser.IndexOnlyContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#Func0Exp.
|
||||
def visitFunc0Exp(self, ctx:MathExprParser.Func0ExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -422,6 +422,87 @@ class MathExprListener(ParseTreeListener):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceWithStep.
|
||||
def enterSliceWithStep(self, ctx:MathExprParser.SliceWithStepContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceWithStep.
|
||||
def exitSliceWithStep(self, ctx:MathExprParser.SliceWithStepContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#Slice.
|
||||
def enterSlice(self, ctx:MathExprParser.SliceContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#Slice.
|
||||
def exitSlice(self, ctx:MathExprParser.SliceContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceColonStartStep.
|
||||
def enterSliceColonStartStep(self, ctx:MathExprParser.SliceColonStartStepContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceColonStartStep.
|
||||
def exitSliceColonStartStep(self, ctx:MathExprParser.SliceColonStartStepContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceColonStart.
|
||||
def enterSliceColonStart(self, ctx:MathExprParser.SliceColonStartContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceColonStart.
|
||||
def exitSliceColonStart(self, ctx:MathExprParser.SliceColonStartContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceExprColonStep.
|
||||
def enterSliceExprColonStep(self, ctx:MathExprParser.SliceExprColonStepContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceExprColonStep.
|
||||
def exitSliceExprColonStep(self, ctx:MathExprParser.SliceExprColonStepContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceExpr.
|
||||
def enterSliceExpr(self, ctx:MathExprParser.SliceExprContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceExpr.
|
||||
def exitSliceExpr(self, ctx:MathExprParser.SliceExprContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceColonStep.
|
||||
def enterSliceColonStep(self, ctx:MathExprParser.SliceColonStepContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceColonStep.
|
||||
def exitSliceColonStep(self, ctx:MathExprParser.SliceColonStepContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SliceOnly.
|
||||
def enterSliceOnly(self, ctx:MathExprParser.SliceOnlyContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SliceOnly.
|
||||
def exitSliceOnly(self, ctx:MathExprParser.SliceOnlyContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#IndexOnly.
|
||||
def enterIndexOnly(self, ctx:MathExprParser.IndexOnlyContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#IndexOnly.
|
||||
def exitIndexOnly(self, ctx:MathExprParser.IndexOnlyContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#Func0Exp.
|
||||
def enterFunc0Exp(self, ctx:MathExprParser.Func0ExpContext):
|
||||
pass
|
||||
|
||||
+3412
-3002
File diff suppressed because one or more lines are too long
@@ -239,6 +239,51 @@ class MathExprVisitor(ParseTreeVisitor):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceWithStep.
|
||||
def visitSliceWithStep(self, ctx:MathExprParser.SliceWithStepContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#Slice.
|
||||
def visitSlice(self, ctx:MathExprParser.SliceContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceColonStartStep.
|
||||
def visitSliceColonStartStep(self, ctx:MathExprParser.SliceColonStartStepContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceColonStart.
|
||||
def visitSliceColonStart(self, ctx:MathExprParser.SliceColonStartContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceExprColonStep.
|
||||
def visitSliceExprColonStep(self, ctx:MathExprParser.SliceExprColonStepContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceExpr.
|
||||
def visitSliceExpr(self, ctx:MathExprParser.SliceExprContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceColonStep.
|
||||
def visitSliceColonStep(self, ctx:MathExprParser.SliceColonStepContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SliceOnly.
|
||||
def visitSliceOnly(self, ctx:MathExprParser.SliceOnlyContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#IndexOnly.
|
||||
def visitIndexOnly(self, ctx:MathExprParser.IndexOnlyContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#Func0Exp.
|
||||
def visitFunc0Exp(self, ctx:MathExprParser.Func0ExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
@@ -0,0 +1,191 @@
|
||||
"""Tests for tensor slicing functionality."""
|
||||
import sys
|
||||
import os
|
||||
import torch
|
||||
import pytest
|
||||
|
||||
# Ensure we can import the module
|
||||
_here = os.path.abspath(os.path.dirname(__file__))
|
||||
_project_root = os.path.abspath(os.path.join(_here, os.pardir))
|
||||
if _project_root not in sys.path:
|
||||
sys.path.insert(0, _project_root)
|
||||
|
||||
from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
|
||||
from more_math.Parser.MathExprLexer import MathExprLexer
|
||||
from more_math.Parser.MathExprParser import MathExprParser
|
||||
from antlr4 import InputStream, CommonTokenStream
|
||||
|
||||
|
||||
def parse_and_visit(expr_str, variables):
|
||||
"""Helper to parse and visit an expression."""
|
||||
lexer = MathExprLexer(InputStream(expr_str))
|
||||
stream = CommonTokenStream(lexer)
|
||||
parser = MathExprParser(stream)
|
||||
tree = parser.start()
|
||||
|
||||
shape = (1, 1, 1, 1)
|
||||
visitor = UnifiedMathVisitor(variables, shape)
|
||||
return visitor.visit(tree)
|
||||
|
||||
|
||||
class TestTensorSlicing:
|
||||
"""Test tensor slicing operations."""
|
||||
|
||||
def test_basic_slice_1d(self):
|
||||
"""Test basic 1D tensor slicing."""
|
||||
t = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
|
||||
vars = {"t": t}
|
||||
|
||||
# Test [1:3]
|
||||
result = parse_and_visit("t[1:3]", vars)
|
||||
assert torch.allclose(result, torch.tensor([2.0, 3.0]))
|
||||
|
||||
# Test [0:2]
|
||||
result = parse_and_visit("t[0:2]", vars)
|
||||
assert torch.allclose(result, torch.tensor([1.0, 2.0]))
|
||||
|
||||
# Test [2:]
|
||||
result = parse_and_visit("t[2:]", vars)
|
||||
assert torch.allclose(result, torch.tensor([3.0, 4.0, 5.0]))
|
||||
|
||||
# Test [:3]
|
||||
result = parse_and_visit("t[:3]", vars)
|
||||
assert torch.allclose(result, torch.tensor([1.0, 2.0, 3.0]))
|
||||
|
||||
def test_slice_with_step_1d(self):
|
||||
"""Test 1D tensor slicing with step."""
|
||||
t = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
|
||||
vars = {"t": t}
|
||||
|
||||
# Test [::2]
|
||||
result = parse_and_visit("t[::2]", vars)
|
||||
assert torch.allclose(result, torch.tensor([1.0, 3.0, 5.0]))
|
||||
|
||||
# Test [1::2]
|
||||
result = parse_and_visit("t[1::2]", vars)
|
||||
assert torch.allclose(result, torch.tensor([2.0, 4.0, 6.0]))
|
||||
|
||||
# Test [0:4:2]
|
||||
result = parse_and_visit("t[0:4:2]", vars)
|
||||
assert torch.allclose(result, torch.tensor([1.0, 3.0]))
|
||||
|
||||
def test_slice_2d(self):
|
||||
"""Test 2D tensor slicing."""
|
||||
t = torch.tensor([
|
||||
[1.0, 2.0, 3.0],
|
||||
[4.0, 5.0, 6.0],
|
||||
[7.0, 8.0, 9.0]
|
||||
])
|
||||
vars = {"t": t}
|
||||
|
||||
# Test [0:2, 1:3]
|
||||
result = parse_and_visit("t[0:2, 1:3]", vars)
|
||||
expected = torch.tensor([[2.0, 3.0], [5.0, 6.0]])
|
||||
assert torch.allclose(result, expected)
|
||||
|
||||
# Test [1:, :2]
|
||||
result = parse_and_visit("t[1:, :2]", vars)
|
||||
expected = torch.tensor([[4.0, 5.0], [7.0, 8.0]])
|
||||
assert torch.allclose(result, expected)
|
||||
|
||||
def test_negative_indices(self):
|
||||
"""Test slicing with negative indices."""
|
||||
t = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
|
||||
vars = {"t": t}
|
||||
|
||||
# Test [-2:]
|
||||
result = parse_and_visit("t[-2:]", vars)
|
||||
assert torch.allclose(result, torch.tensor([4.0, 5.0]))
|
||||
|
||||
# Test [:-2]
|
||||
result = parse_and_visit("t[:-2]", vars)
|
||||
assert torch.allclose(result, torch.tensor([1.0, 2.0, 3.0]))
|
||||
|
||||
def test_indexing_still_works(self):
|
||||
"""Test that regular indexing still works after slicing changes."""
|
||||
t = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
|
||||
vars = {"t": t}
|
||||
|
||||
# Test [0]
|
||||
result = parse_and_visit("t[0]", vars)
|
||||
assert result == 1.0
|
||||
|
||||
# Test [2]
|
||||
result = parse_and_visit("t[2]", vars)
|
||||
assert result == 3.0
|
||||
|
||||
# Test [-1]
|
||||
result = parse_and_visit("t[-1]", vars)
|
||||
assert result == 5.0
|
||||
|
||||
def test_list_slicing(self):
|
||||
"""Test list slicing."""
|
||||
l = [1.0, 2.0, 3.0, 4.0, 5.0]
|
||||
vars = {"l": l}
|
||||
|
||||
# Test [1:3]
|
||||
result = parse_and_visit("l[1:3]", vars)
|
||||
assert result == [2.0, 3.0]
|
||||
|
||||
# Test [::2]
|
||||
result = parse_and_visit("l[::2]", vars)
|
||||
assert result == [1.0, 3.0, 5.0]
|
||||
|
||||
# Test [-2:]
|
||||
result = parse_and_visit("l[-2:]", vars)
|
||||
assert result == [4.0, 5.0]
|
||||
|
||||
def test_string_slicing(self):
|
||||
"""Test string slicing."""
|
||||
s = "hello"
|
||||
vars = {"s": s}
|
||||
|
||||
# Test [1:3]
|
||||
result = parse_and_visit("s[1:3]", vars)
|
||||
assert result == "el"
|
||||
|
||||
# Test [::2]
|
||||
result = parse_and_visit("s[::2]", vars)
|
||||
assert result == "hlo"
|
||||
|
||||
# Test [2:]
|
||||
result = parse_and_visit("s[2:]", vars)
|
||||
assert result == "llo"
|
||||
|
||||
|
||||
class TestListAndTensorIndexing:
|
||||
"""Test that existing indexing functionality still works."""
|
||||
|
||||
def test_tensor_indexing_multi_dim(self):
|
||||
"""Test multi-dimensional tensor indexing."""
|
||||
t = torch.arange(24).reshape(2, 3, 4).float()
|
||||
vars = {"t": t}
|
||||
|
||||
# Test [0, 1, 2]
|
||||
result = parse_and_visit("t[0, 1, 2]", vars)
|
||||
assert result == 6.0
|
||||
|
||||
# Test [1, 2, 3]
|
||||
result = parse_and_visit("t[1, 2, 3]", vars)
|
||||
assert result == 23.0
|
||||
|
||||
def test_list_indexing(self):
|
||||
"""Test list indexing."""
|
||||
l = [10, 20, 30, 40, 50]
|
||||
vars = {"l": l}
|
||||
|
||||
# Test [0]
|
||||
result = parse_and_visit("l[0]", vars)
|
||||
assert result == 10
|
||||
|
||||
# Test [3]
|
||||
result = parse_and_visit("l[3]", vars)
|
||||
assert result == 40
|
||||
|
||||
# Test [-1]
|
||||
result = parse_and_visit("l[-1]", vars)
|
||||
assert result == 50
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Reference in New Issue
Block a user