feat: extend EpochMetric to support more compute_fn output types

[zongyang078]

Fixes #1757

Description:
Extends EpochMetric to support more output types returned by compute_fn, as requested in #1757.

Previously, compute_fn could only return scalar values (int/float). This PR adds support for:

• torch.Tensor
• tuple / list of tensors
• Mapping (dict) with string keys and tensor values

A recursive _check_output_type method validates the output, including nested containers. Unsupported types raise a clear TypeError.

Distributed handling (world_size > 1):

• Scalar and torch.Tensor outputs are broadcast across processes
• tuple/list/Mapping outputs raise NotImplementedError (single-process only for now)
• A status-code broadcast ensures all ranks fail symmetrically without deadlocking

Tests:

• Tensor, tuple, list, dict outputs (with value verification)
• Unsupported type (str) raises TypeError
• Nested invalid type raises TypeError
• Mapping with non-str key raises TypeError

Note: This iteration adopts a more conservative approach to distributed handling, with recursive type validation and NotImplementedError for non-scalar container outputs when world_size > 1.

Check list:

• New tests are added
• New doc strings: description and/or example code are in RST format
• Documentation is updated (docstring reflects new supported types and distributed behavior)

[aaishwarymishra]

Just to be sure we also should check if the output values are as expected or correct after the computation.

[zongyang078]

Added value verification for all output types — each test now checks that the result matches directly calling compute_fn on the concatenated data.

[Abhisri436]

Hi! I’m working on issue #1757 for a course contribution and noticed this PR already covers the main source changes.

I added value-checking tests on my fork for scalar tensor, vector tensor, tuple/list of tensors, and mapping outputs. These tests check the actual computed values, which seems related to the review comment about verifying output correctness.

Test commit: Abhisri436@56345826

Would it be useful for me to adapt these tests into this PR or keep them as a separate reference?

[aaishwarymishra]

@Abhisri436 It's generally not a good Idea to implement the changes as an external pr until explicitly mentioned that the change if a future improvement or can be done separately :)

[aaishwarymishra]

@zongyang078 Are you panning to work on it?

[zongyang078]

@zongyang078 Are you panning to work on it?

Yes, I'm planning to continue! I've reviewed PR #3802 and found several improvements worth incorporating, particularly the recursive _check_output_type method and the conservative distributed handling with NotImplementedError. I'll update this PR with these improvements shortly.

[zongyang078]

The refactor has been pushed. Changes include:

1. Extracted recursive _check_output_type method for nested container validation
2. Conservative distributed handling with NotImplementedError for tuple/list/mapping outputs when world_size > 1
3. Status-code broadcast mechanism to ensure symmetric failure across ranks
4. Added tests for nested invalid types and non-str mapping keys
5. Updated docstring
   All 17 non-distributed tests pass locally.

Also thanks to @Abhisri436 for the discussion in #3802 — the safer distributed approach was a valuable direction to explore.

[Abhisri436]

Hi, thank you for the acknowledgement. I’m glad the safer distributed approach was useful.

Since the latest version of #3789 now incorporates the recursive validation, status-code distributed handling, safe_mode=True broadcast, and conservative NotImplementedError behavior first introduced in #3802, I think it would be good to consolidate the two PRs rather than continue with duplicate implementations.

#3802 also includes permanent WORLD_SIZE=2 distributed tests that verify this implementation under real multi-rank execution. These tests remain a distinct part of #3802 and provide coverage that is not currently present in #3789.

Could a maintainer advise on the preferred consolidation path and the appropriate way to preserve attribution for the implementation originally introduced in #3802?

[aaishwarymishra]

Hi @zongyang078 sorry, I missed your previous pushes.
I have few quenstion's

I think the ignite boradcast method itself checks for types in ignite/distributed/comp_models/base.py and will throw error for sequence and mapping even after this check in distributed setting.

[zongyang078]

Hi @zongyang078 sorry, I missed your previous pushes. I have few quenstion's

I think the ignite boradcast method itself checks for types in ignite/distributed/comp_models/base.py and will throw error for sequence and mapping even after this check in distributed setting.

Hi @aaishwarymishra, you're right that broadcast() in base.py rejects Sequence/Mapping. The status-code exists because only rank 0 runs compute_fn — without it, rank 0 would hit that TypeError and exit while other ranks deadlock waiting in broadcast's safe_mode path.

That said, I think the better direction is to support container outputs in distributed mode by recursively broadcasting each leaf tensor individually. Would that be welcome?

[aaishwarymishra]

@vfdev-5 what's your thoughts on it?