Skip to content

Fix slogdet specialization for batched input - #2411

Open
N-0-MAD wants to merge 5 commits into
pymc-devs:mainfrom
N-0-MAD:fix-batched-slogdet
Open

N-0-MAD wants to merge 5 commits into
pymc-devs:mainfrom
N-0-MAD:fix-batched-slogdet

Conversation

@N-0-MAD

@N-0-MAD N-0-MAD commented Sep 13, 2026

Copy link
Copy Markdown

Description

Fix slogdet_specialization for batched inputs.

The slogdet_specialization rewrite currently replaces a batched
Det operation with a bare SLogDet. Since SLogDet only accepts
2D matrix inputs, this causes the rewrite to fail for batched inputs.

This change wraps SLogDet in Blockwise, preserving the batch
dimensions while applying SLogDet to each matrix in the batch.

Added regression coverage for batched:

  • sign(det(x))
  • log(abs(det(x)))
  • log(det(x))

Related Issue

Checklist

Type of change

  • Bug fix
  • New feature / enhancement
  • Documentation
  • Maintenance
  • Other (please specify):

nodes = f.maker.fgraph.apply_nodes
assert not any(isinstance(node.op, SLogDet) for node in nodes)

# Batched input

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

can you instead parameterize the existing test with batched dim. and refactor it to use RewriteTester helper (called something like that but I'm never sure of the name)

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Thanks! I've parameterized the existing test over batch_dim and refactored it to use RewriteTester as suggested. I've also removed the separate batched-input test since it's now covered by the parameterized test.

rtol=1e-3 if config.floatX == "float32" else 1e-8,
)
result = RewriteTester([x], [log_det_x], include=["stabilize", "specialize"])
result.assert_eval(a)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

why no eval graph here?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

I had added the expected graph here, but assert_graph fails for the batched case because the rewritten graph represents the scalar constants as broadcast vectors ([-1]/[nan]), while the manually constructed expected graph keeps them as ExpandDims of scalars. The numerical evaluation passes for both cases, so I kept assert_eval for this one.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

The manually constructed graph can be manually constructed as specific as needed?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Done. Updated the expected graph to explicitly construct the scalar vs. broadcast-vector constants for the two cases. assert_graph now passes for both batched and non-batched inputs.

atol=1e-3 if config.floatX == "float32" else 1e-8,
rtol=1e-3 if config.floatX == "float32" else 1e-8,
)
_, log_abs_x = Blockwise(SLogDet())(x)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

this variable name is confusing (not det?)

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Yes, thanks! Renamed it to expected_log_abs_det_x for clarity.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

slogdet cannot be built for a batched input on any backend

2 participants