Skip to content

Fixes for scaled initializers, function inlining, model selection, and added Conv1d support. #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sidcraftscode wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fixes for scaled initializers, function inlining, model selection, and added Conv1d support. #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions delta/backend/fx_lowering.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -527,6 +527,8 @@ def _lower_layer(self, node: Layer, ctx: FXContext) -> Node:
# Instantiate the actual torch module
if node.kind == 'Linear':
self.layers[layer_id] = torch.nn.Linear(*node.args, **node.kwargs)
elif node.kind == 'Conv1d':
self.layers[layer_id] = torch.nn.Conv1d(*node.args, **node.kwargs)
elif node.kind == 'Conv2d':
self.layers[layer_id] = torch.nn.Conv2d(*node.args, **node.kwargs)
elif node.kind == 'Embedding':
Expand Down
15 changes: 5 additions & 10 deletions delta/ir/sir_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,9 +256,7 @@ def _build_obs_decl(self, stmt: ObsDecl) -> ObsRef:

def _build_constraint_stmt(self, stmt: ASTConstraintStmt) -> ConstraintOp:
"""Build a constraint statement."""
print(f"DEBUG: Building constraint stmt: {stmt}")
expr_node = self._build_expr(stmt.expr)
print(f"DEBUG: Constraint expr_node created: {type(expr_node)}")

# Determine constraint kind
if stmt.kind == ASTConstraintKind.REQUIRE:
Expand All @@ -279,18 +277,16 @@ def _build_constraint_stmt(self, stmt: ASTConstraintStmt) -> ConstraintOp:
if stmt.slack:
slack = self._build_expr(stmt.slack)

print("DEBUG: Before SIRProperty")
props = SIRProperty(
dtype=FloatType(),
requires_grad=expr_node.requires_grad,
dtype=ConstraintType(),
Copy link

Copilot AI Jan 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The type ConstraintType() does not exist in the codebase. This will cause a NameError at runtime when building constraint statements. The type should be FloatType() instead, as constraints evaluate to scalar penalty terms that are added to the objective. The original code used FloatType() for the dtype of constraints.

Suggested change
dtype=ConstraintType(),
dtype=FloatType(),

Copilot uses AI. Check for mistakes.
role=RoleInfo.constraint(),
location=stmt.location
)

print("DEBUG: Before ConstraintOp")
return ConstraintOp(
kind=kind,
lhs=expr_node,
rhs=None,
rhs=None, # All constraints normalized to expr checks
weight=weight,
slack=slack,
_props=props
Expand Down Expand Up @@ -358,7 +354,6 @@ def _build_if_stmt(self, stmt: IfStmt) -> Optional[SIRNode]:
gate = GateOp(compare=TensorOpKind.GT, lhs=condition, rhs=Const(value=0.0), temperature=temperature)
then_val = then_block.result or Const(value=0.0)
else_val = (else_block.result if else_block else None) or Const(value=0.0)
print(f"DEBUG: IfStmt temperature={temperature}, then_result={then_block.result}, else_result={else_block.result if else_block else 'None'}")
return MixOp(gate, then_val, else_val)

return None
Expand Down Expand Up @@ -447,9 +442,9 @@ def _build_block(self, block: Block) -> SIRBlock:
if block.result:
result = self._build_expr(block.result)
nodes.append(result)
elif nodes and not isinstance(block.statements[-1], (ReturnStmt, ParamDecl, ObsDecl)):
elif nodes and not isinstance(block.statements[-1], (ParamDecl, ObsDecl)):
Copy link

Copilot AI Jan 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The removal of ReturnStmt from this check may cause incorrect block result inference. When a block ends with a return statement, the return statement itself should not be used as the block's implicit result value. Consider whether this change is intentional and whether it aligns with the semantics of how return statements should be handled in blocks.

Copilot uses AI. Check for mistakes.
# Fallback: if no explicit result, use the last node as result (like Rust/Ruby)
# but only if it's not a return or declaration
# Declared params/obs don't produce values
result = nodes[-1]

return SIRBlock(nodes=nodes, result=result)
Expand Down
18 changes: 15 additions & 3 deletions delta/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
# Note: We import frontend AST types here for initializer evaluation.
# Ideally this metadata would be extracted during compilation, but
# for now we use the AST directly since CompileResult exposes it.
from delta.frontend.ast import ParamDecl, Call, Identifier, Literal
from delta.frontend.ast import ParamDecl, Call, Identifier, Literal, BinaryOp
Copy link

Copilot AI Jan 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Import of 'BinaryOp' is not used.

Suggested change
from delta.frontend.ast import ParamDecl, Call, Identifier, Literal, BinaryOp
from delta.frontend.ast import ParamDecl, Call, Identifier, Literal

Copilot uses AI. Check for mistakes.


@dataclass
Expand All @@ -48,7 +48,11 @@ class DeltaModel:

def __init__(self, compile_result: CompileResult):
self._compile_result = compile_result
self._module = list(compile_result.graph_modules.values())[0]
# Prefer 'forward' module if available (standard for training), otherwise take the first one
if compile_result.graph_modules and "forward" in compile_result.graph_modules:
self._module = compile_result.graph_modules["forward"]
else:
self._module = list(compile_result.graph_modules.values())[0]
self._device = torch.device("cpu")
self._init_params_from_sir()

Expand Down Expand Up @@ -150,7 +154,15 @@ def _init_from_name(self, name: str, shape: tuple) -> torch.Tensor:

def _extract_shape_from_initializer(self, initializer_expr: Any) -> Optional[tuple]:
"""Extract shape from an initializer expression like randn(3, 2)."""
from delta.frontend.ast import Call, Identifier, Literal, Tensor
from delta.frontend.ast import Call, Identifier, Literal, Tensor, BinaryOp
Copy link

Copilot AI Jan 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The BinaryOp import was added to the top-level imports but is also imported again within the _extract_shape_from_initializer method at line 157. Consider removing the redundant local import since BinaryOp is now available from the module-level import at line 30.

Suggested change
from delta.frontend.ast import Call, Identifier, Literal, Tensor, BinaryOp
from delta.frontend.ast import Tensor

Copilot uses AI. Check for mistakes.

# Handle scaling: randn(...) * 0.01
if isinstance(initializer_expr, BinaryOp):
# Recursively check operands
shape = self._extract_shape_from_initializer(initializer_expr.left)
if shape: return shape
return self._extract_shape_from_initializer(initializer_expr.right)
Comment on lines +159 to +164
Copy link

Copilot AI Jan 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The handling of BinaryOp expressions in _extract_shape_from_initializer is incomplete. While it recursively extracts shapes from binary operations like randn(3, 2) * 0.01, the corresponding _eval_initializer method (lines 90-137) doesn't handle BinaryOp expressions. This means scaled initializers will have their shapes extracted correctly, but the actual initialization won't apply the scaling factor. Consider also updating _eval_initializer to handle BinaryOp expressions to fully support scaled initializers.

Copilot uses AI. Check for mistakes.

if isinstance(initializer_expr, Call) and isinstance(initializer_expr.func, Identifier):
func_name = initializer_expr.func.name
args = initializer_expr.args
Expand Down
Loading