Conversation
| nodes = f.maker.fgraph.apply_nodes | ||
| assert not any(isinstance(node.op, SLogDet) for node in nodes) | ||
|
|
||
| # Batched input |
There was a problem hiding this comment.
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)
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
The manually constructed graph can be manually constructed as specific as needed?
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
this variable name is confusing (not det?)
There was a problem hiding this comment.
Yes, thanks! Renamed it to expected_log_abs_det_x for clarity.
Description
Fix
slogdet_specializationfor batched inputs.The
slogdet_specializationrewrite currently replaces a batchedDetoperation with a bareSLogDet. SinceSLogDetonly accepts2D matrix inputs, this causes the rewrite to fail for batched inputs.
This change wraps
SLogDetinBlockwise, preserving the batchdimensions while applying
SLogDetto 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