feat: Extend BatchNorm with rank-matched channel tensors and optional momentum

- **Rank-matched channel tensors**: SCALE, BIAS, MEAN, VAR, SAVED_MEAN, and
  SAVED_INV_VARIANCE now accept tensors of shape [1, C, 1, ..., 1] (ones in
  all non-feature dimensions) in addition to the canonical 1D [C] shape.
  The node validates and infers strides for both forms.  The ASM emitter
  collapses rank-matched inputs to 1D via `flatten.using_ints` before
  `native_batch_norm`, and expands rank-matched stat outputs back via
  `reshape` after.

- **Optional momentum**: Momentum is no longer required.  When omitted, the
  ASM emitter emits `torch.constant.float 1.000000e-01` (PyTorch default).
  If provided it must still be a scalar constant.

- **Tests**:
  - `test_batchnorm_node.cpp`: updated "Momentum missing" section to expect
    success; added rank-matched validation sections.
  - `test_batchnorm_infer_asm_emitter_nchw_rank_matched.cpp`: lit test
    verifying `flatten.using_ints` collapse for rank-matched MEAN/VAR.
  - `test_batchnorm_infer_asm_emitter_nchw_no_momentum.cpp`: lit test
    verifying `torch.constant.float` default when momentum is absent.

- **Samples**:
  - `batchnorm_infer_nchw_rank_matched_scale_bias.cpp`: end-to-end inference
    with [1,C,1,1] scale, bias, mean, var.
  - `batchnorm_infer_nchw_no_momentum.cpp`: end-to-end inference with no
    momentum tensor.
  - `batchnorm_infer_ncdhw_rank_matched.cpp`: end-to-end 5D (NCDHW) inference
    with rank-matched [1,C,1,1,1] channel tensors.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Rob Suderman <rob.suderman@gmail.com>

Author: rsuderman

include/fusilli/node/batchnorm_node.h

@@ -125,12 +125,12 @@ class BatchNormNode : public NodeCRTP<BatchNormNode> {
     FUSILLI_RETURN_ERROR_IF(!eT->isScalar(), ErrorCode::InvalidAttribute,
                             "BatchNorm epsilon must be a scalar constant");
 
-    // Momentum checks.
+    // Momentum checks (optional — omitting uses the PyTorch default 0.1).
     std::shared_ptr<TensorAttr> mT = batchnormAttr.getMomentum();
-    FUSILLI_RETURN_ERROR_IF(!mT, ErrorCode::AttributeNotSet,
-                            "BatchNorm momentum not set");
-    FUSILLI_RETURN_ERROR_IF(!mT->isScalar(), ErrorCode::InvalidAttribute,
-                            "BatchNorm momentum must be a scalar constant");
+    if (mT) {
+      FUSILLI_RETURN_ERROR_IF(!mT->isScalar(), ErrorCode::InvalidAttribute,
+                              "BatchNorm momentum must be a scalar constant");
+    }
 
     return ok();
   }
@@ -151,8 +151,19 @@ class BatchNormNode : public NodeCRTP<BatchNormNode> {
     auto infer1DTensor = [&](const std::shared_ptr<TensorAttr> &t) {
       if (t->getDim().empty())
         t->setDim(channel1DDim);
-      if (t->getStride().empty())
-        t->setStride(channel1DStride);
+      if (t->getStride().empty()) {
+        const std::vector<int64_t> &dim = t->getDim();
+        size_t rank = dim.size();
+        if (rank == xDim.size()) {
+          // Rank-matched tensor [1, C, 1, ...]: compute contiguous stride.
+          std::vector<int64_t> stride(rank, 1);
+          for (int64_t i = (int64_t)rank - 2; i >= 0; --i)
+            stride[i] = stride[i + 1] * dim[i + 1];
+          t->setStride(stride);
+        } else {
+          t->setStride(channel1DStride);
+        }
+      }
     };
 
     // Infer 1D channel tensors.
@@ -203,18 +214,49 @@ class BatchNormNode : public NodeCRTP<BatchNormNode> {
                                 "defined by its stride");
 
     // Shape checks for 1D channel tensors.
+    // Accepts either the canonical 1D form [C] with unit stride, or a
+    // rank-matched form [1, C, 1, ..., 1] with unit stride at the channel dim.
     auto check1DShape = [&](const std::shared_ptr<TensorAttr> &t,
                             const std::string &name) -> ErrorObject {
       if (!t)
         return ok();
+      const std::vector<int64_t> &tDim = t->getDim();
+      const std::vector<int64_t> &tStride = t->getStride();
+      size_t xRank = xDim.size();
+
+      if (tDim == expectedCDim) {
+        // Canonical 1D form [C]: stride must be {1}.
+        FUSILLI_RETURN_ERROR_IF(
+            tStride != std::vector<int64_t>{1}, ErrorCode::InvalidAttribute,
+            "BatchNorm tensor " + name + " must have unit stride");
+        return ok();
+      }
+
+      if (tDim.size() == xRank) {
+        // Rank-matched form [1, C, 1, ..., 1]: channel dim must equal C and
+        // all other dims must be 1.
+        bool validShape = (tDim[1] == xDim[1]);
+        for (size_t i = 0; i < xRank && validShape; ++i)
+          if (i != 1 && tDim[i] != 1)
+            validShape = false;
+        FUSILLI_RETURN_ERROR_IF(
+            !validShape, ErrorCode::InvalidAttribute,
+            "BatchNorm tensor " + name +
+                " must be 1D with size equal to channel dimension C"
+                " or rank-matched with ones in all non-feature dimensions");
+        FUSILLI_RETURN_ERROR_IF(
+            tStride.size() != xRank || tStride[1] != 1,
+            ErrorCode::InvalidAttribute,
+            "BatchNorm tensor " + name +
+                " must have unit stride at the channel dimension");
+        return ok();
+      }
+
       FUSILLI_RETURN_ERROR_IF(
-          t->getDim() != expectedCDim, ErrorCode::InvalidAttribute,
+          true, ErrorCode::InvalidAttribute,
           "BatchNorm tensor " + name +
-              " must be 1D with size equal to channel dimension C");
-      FUSILLI_RETURN_ERROR_IF(t->getStride() != std::vector<int64_t>{1},
-                              ErrorCode::InvalidAttribute,
-                              "BatchNorm tensor " + name +
-                                  " must have unit stride");
+              " must be 1D with size equal to channel dimension C"
+              " or rank-matched with ones in all non-feature dimensions");
       return ok();
     };