onnx-ir 0.21.0

//! # Gather
//!
//! Gathers elements from input tensor along a specified axis using indices.
//!
//! **ONNX Spec**: <https://onnx.ai/onnx/operators/onnx__Gather.html>
//!
//! ## Description
//!
//! Given `data` tensor of rank r >= 1, and `indices` tensor of rank q, gather entries of the
//! axis dimension of `data` (by default outer-most one as axis=0) indexed by `indices`, and
//! concatenates them in an output tensor of rank q + (r - 1).
//!
//! ## Type Constraints
//!
//! - T: tensor(bfloat16), tensor(bool), tensor(complex128), tensor(complex64), tensor(double),
//!   tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8),
//!   tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
//! - Tind: tensor(int32), tensor(int64)
//!
//! ## Opset Versions
//!
//! - **Opset 1**: Initial version with basic gather functionality.
//! - **Opset 11**: Added support for negative indices; out-of-bounds indices now raise an error instead of undefined behavior.
//! - **Opset 13**: Added bfloat16 type support; no functional changes to operation semantics.
//!
//! **Implementation Note**: This implementation validates opset 11+ (see FIXME at line 92).
use derive_new::new;
use onnx_ir_derive::NodeBuilder;

use crate::ir::{ArgType, Argument, Node, RawNode, TensorType};
use crate::processor::{
    InputPreferences, InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec,
    ProcessError,
};

/// Configuration for the Gather operation.
#[derive(Debug, Clone, new)]
pub struct GatherConfig {
    pub axis: usize,
}

/// Node representation for Gather operation
#[derive(Debug, Clone, NodeBuilder)]
pub struct GatherNode {
    pub name: String,
    pub inputs: Vec<Argument>,
    pub outputs: Vec<Argument>,
    pub config: GatherConfig,
}

pub(crate) struct GatherProcessor;

impl NodeProcessor for GatherProcessor {
    type Config = GatherConfig;

    fn spec(&self) -> NodeSpec {
        NodeSpec {
            min_opset: 1,
            max_opset: None,
            inputs: InputSpec::Exact(2),
            outputs: OutputSpec::Exact(1),
        }
    }

    fn is_noop(&self, node: &RawNode) -> bool {
        // Gather is a no-op when input is Scalar (single element, nothing to gather)
        matches!(node.inputs[0].ty, ArgType::ScalarNative(_))
    }

    fn input_preferences(
        &self,
        node: &RawNode,
        _opset: usize,
    ) -> Result<Option<InputPreferences>, ProcessError> {
        use crate::processor::ArgPreference;

        if node.inputs.len() < 2 {
            return Ok(None);
        }

        // When gathering from Shape data, prefer scalar indices as ScalarNative
        // (array indexing uses native values) and multi-element indices as Shape
        if node.inputs[0].ty.is_shape() {
            let pref = if node.inputs[1].ty.is_scalar() {
                ArgPreference::ScalarNative
            } else {
                ArgPreference::Shape
            };
            Ok(Some(
                InputPreferences::new().add(&node.inputs[1].name, pref),
            ))
        } else {
            Ok(None)
        }
    }

    fn infer_types(
        &self,
        node: &mut RawNode,
        _opset: usize,
        _output_preferences: &OutputPreferences,
    ) -> Result<(), ProcessError> {
        // TODO: Validate indices tensor type is int32 or int64 per ONNX spec - Missing type constraint validation

        // Extract the input rank for axis normalization
        // Scalar inputs have effective rank 1 (single element that can be gathered)
        let input_dim = match &node.inputs[0].ty {
            ArgType::Tensor(tensor) => tensor.rank as i64,
            ArgType::Shape(shape_rank) => *shape_rank as i64,
            ArgType::ScalarTensor(_) | ArgType::ScalarNative(_) => 1,
        };

        // Extract the axis attribute (default: 0 per ONNX spec)
        let mut axis: i64 = 0;
        for (key, value) in node.attrs.iter() {
            match key.as_str() {
                "axis" => axis = value.clone().into_i64(),
                _ => {
                    return Err(ProcessError::InvalidAttribute {
                        name: key.clone(),
                        reason: format!("Unexpected attribute for Gather: {}", key),
                    });
                }
            }
        }
        // TODO: Validate negative indices support for opset < 11 - Negative indices added in opset 11, should error for earlier opsets - Missing opset-specific validation

        // Normalize negative axis
        if axis < 0 {
            axis += input_dim;
        }

        // Validate axis is within bounds
        if axis < 0 || axis >= input_dim {
            return Err(ProcessError::InvalidAttribute {
                name: "axis".to_string(),
                reason: format!("axis {} is out of bounds for rank {}", axis, input_dim),
            });
        }

        // Infer output type based on indices rank.
        // Scalar indices (both ScalarTensor and ScalarNative) are rank 0 per ONNX spec.
        let indices_rank = match &node.inputs[1].ty {
            ArgType::Tensor(tensor) => tensor.rank,
            // Scalar indices (both ScalarTensor and ScalarNative) are rank 0,
            // reducing output rank by 1 per ONNX spec: output_rank = q + (r - 1)
            ArgType::ScalarTensor(_) | ArgType::ScalarNative(_) => 0,
            ArgType::Shape(_shape_rank) => {
                1 // Shape indices are always treated as rank 1 for gather
            }
        };

        match &node.inputs[0].ty {
            ArgType::Tensor(input_tensor) => {
                // Output of rank q+(r-1), where q is rank of indices tensor and r is rank of input
                let output_rank = indices_rank + input_tensor.rank - 1;

                if output_rank == 0 {
                    // Output is scalar tensor (stays on device)
                    // Downstream consumers that need native will request ScalarNative via preferences
                    node.outputs[0].ty = ArgType::ScalarTensor(input_tensor.dtype);
                } else {
                    // Output is tensor
                    node.outputs[0].ty = ArgType::Tensor(TensorType {
                        dtype: input_tensor.dtype,
                        rank: output_rank,
                        static_shape: None,
                    });
                }
            }
            ArgType::Shape(_shape_rank) => {
                // When gathering from a shape:
                // - If indices are scalar (rank 0), output is a scalar (single dimension value)
                // - Otherwise, output is a shape with same dimension as indices
                if indices_rank == 0 {
                    node.outputs[0].ty = ArgType::ScalarNative(crate::ir::DType::I64);
                } else {
                    // Shape(N) means N elements, so use the number of gathered elements
                    let output_shape_rank = match &node.inputs[1].ty {
                        ArgType::Shape(shape_rank) => *shape_rank,
                        ArgType::Tensor(t) => {
                            // Use static shape to get actual element count
                            t.static_shape_known()
                                .and_then(|s| s.first().copied())
                                .unwrap_or(indices_rank)
                        }
                        _ => indices_rank,
                    };
                    node.outputs[0].ty = ArgType::Shape(output_shape_rank);
                }
            }
            ArgType::ScalarTensor(dtype) => {
                node.outputs[0].ty = ArgType::ScalarTensor(*dtype);
            }
            ArgType::ScalarNative(dtype) => {
                node.outputs[0].ty = ArgType::ScalarNative(*dtype);
            }
        }

        // when index is a constant scalar, make it static so that it can be later eliminated
        // should ideally be in lift_constants but at this point we don't know whether the index
        // is scalar
        let is_scalar_index = match &node.inputs[1].ty {
            ArgType::ScalarTensor(_) | ArgType::ScalarNative(_) => true,
            ArgType::Tensor(t) => t.rank == 0,
            _ => false,
        };
        if node.inputs.len() > 1 && node.inputs[1].is_constant() && is_scalar_index {
            node.inputs[1].to_static()?;
        }

        Ok(())
    }

    fn extract_config(&self, node: &RawNode, _opset: usize) -> Result<Self::Config, ProcessError> {
        // Extract the input rank for axis normalization
        // Scalar inputs have effective rank 1 (single element that can be gathered)
        let input_dim = match &node.inputs[0].ty {
            ArgType::Tensor(tensor) => tensor.rank as i64,
            ArgType::Shape(shape_rank) => *shape_rank as i64,
            ArgType::ScalarTensor(_) | ArgType::ScalarNative(_) => 1,
        };

        // Extract the axis attribute (default: 0 per ONNX spec)
        let mut axis: i64 = 0;
        for (key, value) in node.attrs.iter() {
            if key.as_str() == "axis" {
                axis = value.clone().into_i64()
            }
            // TODO: Add validation for unexpected attributes (currently silently ignored)
        }

        // Normalize negative axis
        if axis < 0 {
            axis += input_dim;
        }

        let config = GatherConfig {
            axis: axis as usize,
        };
        Ok(config)
    }

    fn build_node(&self, builder: RawNode, opset: usize) -> Node {
        let config = self
            .extract_config(&builder, opset)
            .expect("Config extraction failed");

        Node::Gather(GatherNode {
            name: builder.name,
            inputs: builder.inputs,
            outputs: builder.outputs,
            config,
        })
    }
}

#[cfg(test)]
mod tests {
    use burn_tensor::DType;

    use super::*;
    use crate::ir::NodeType;
    use crate::node::test_utils::TestNodeBuilder;

    fn create_test_node(axis: i64, input_rank: usize, is_shape: bool) -> TestNodeBuilder {
        // Start building the node with the appropriate input type
        let mut builder =
            TestNodeBuilder::new(NodeType::Gather, "test_gather").attr_int("axis", axis);

        if is_shape {
            builder = builder.add_input("data", ArgType::Shape(input_rank));
        } else {
            builder = builder.input_tensor_f32("data", input_rank, None);
        }

        // Add indices and output
        builder
            .input_tensor_i64("indices", 1, None)
            .output_tensor_f32("output", input_rank, None)
    }

    #[test]
    fn test_gather_config_basic() {
        let node = create_test_node(0, 3, false).build();
        let mut node = node;
        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();
        assert_eq!(config.axis, 0);
    }

    #[test]
    fn test_gather_config_negative_axis() {
        let node = create_test_node(-2, 3, false).build();
        let mut node = node;
        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();
        assert_eq!(config.axis, 1); // -2 + 3 = 1
    }

    #[test]
    fn test_gather_config_shape_input() {
        let node = create_test_node(0, 4, true).build(); // Shape of a 4D tensor
        let mut node = node;
        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();
        assert_eq!(config.axis, 0);
    }

    #[test]
    fn test_gather_config_missing_index() {
        let mut node = create_test_node(0, 3, false).build();
        node.inputs.pop(); // Remove the indices input
        let processor = GatherProcessor;
        let spec = processor.spec();
        let result = crate::processor::validate_node_spec(&node, 16, &spec);
        assert!(matches!(
            result,
            Err(ProcessError::InvalidInputCount {
                expected: 2,
                actual: 1
            })
        ));
    }

    #[test]
    fn test_gather_update_outputs_scalar_result() {
        // Test gather with scalar indices on 1D tensor -> scalar output
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_scalar_gather")
            .attr_int("axis", 0)
            .input_tensor_f32("data", 1, None)
            .add_input("indices", ArgType::ScalarNative(crate::ir::DType::I64))
            .output_tensor_f32("output", 1, None)
            .build();

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        let _config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        // Should output ScalarTensor (stays on device)
        match &node.outputs[0].ty {
            ArgType::ScalarTensor(elem_type) => {
                assert_eq!(*elem_type, crate::ir::DType::F32);
            }
            other => panic!("Expected ScalarTensor output, got {:?}", other),
        }
    }

    #[test]
    fn test_gather_update_outputs_tensor_result() {
        // Test gather with 1D indices on 2D tensor -> 2D tensor output
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_tensor_gather")
            .attr_int("axis", 0)
            .input_tensor_f32("data", 2, None)
            .input_tensor_i64("indices", 1, None)
            .output_tensor_f32("output", 2, None)
            .build();

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        let _config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        // Should output tensor with rank 2 (1 + 2 - 1)
        match &node.outputs[0].ty {
            ArgType::Tensor(tensor) => {
                assert_eq!(tensor.rank, 2);
                assert_eq!(tensor.dtype, crate::ir::DType::F32);
            }
            other => panic!("Expected tensor output, got {:?}", other),
        }
    }

    #[test]
    fn test_gather_update_outputs_shape_indices() {
        // Test gather with Shape indices - this was the bug that caused the original issue
        // Gathering from a shape tensor using shape indices should work correctly
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_gather_shape_indices")
            .attr_int("axis", 0)
            .input_shape("data", 3) // Shape input (represents shape of a 3D tensor)
            .add_input("indices", ArgType::Shape(1)) // Shape(1) indices - this was causing the panic
            .output_shape("output", 1) // Output should be Shape(1)
            .build();

        // This should not panic - it was panicking before the fix
        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        let _config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        // Should output Shape(1) since we're gathering from Shape(3) with Shape(1) indices
        match &node.outputs[0].ty {
            ArgType::Shape(rank) => {
                assert_eq!(*rank, 1);
            }
            other => panic!("Expected Shape output, got {:?}", other),
        }
    }

    #[test]
    fn test_gather_update_outputs_shape_scalar_indices() {
        // Test gather with scalar indices on shape input -> scalar output
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_gather_shape_scalar")
            .attr_int("axis", 0)
            .input_shape("data", 2) // Shape input (represents shape of a 2D tensor)
            .add_input("indices", ArgType::ScalarNative(crate::ir::DType::I64)) // Scalar indices
            .output_tensor_i64("output", 0, None) // Will be updated by gather_update_outputs
            .build();

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        let _config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        // Should output scalar when gathering from shape with scalar indices
        match &node.outputs[0].ty {
            ArgType::ScalarNative(elem_type) => {
                assert_eq!(*elem_type, crate::ir::DType::I64);
            }
            other => panic!("Expected scalar output, got {:?}", other),
        }
    }

    #[test]
    fn test_gather_update_outputs_shape_with_shape_indices_rank_2() {
        // Test gather from Shape with Shape(2) indices -> Shape(2) output
        // This tests our fix where Shape indices preserve their rank in the output
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_gather_shape_shape_2")
            .attr_int("axis", 0)
            .input_shape("data", 4) // Shape input (represents shape of a 4D tensor)
            .add_input("indices", ArgType::Shape(2)) // Shape(2) indices
            .output_shape("output", 1) // Initial output, will be updated
            .build();

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        let _config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        // Should output Shape(2) since indices are Shape(2)
        match &node.outputs[0].ty {
            ArgType::Shape(rank) => {
                assert_eq!(*rank, 2, "Expected Shape(2) output for Shape(2) indices");
            }
            other => panic!("Expected Shape(2) output, got {:?}", other),
        }
    }

    #[test]
    fn test_gather_update_outputs_shape_with_shape_indices_rank_3() {
        // Test gather from Shape with Shape(3) indices -> Shape(3) output
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_gather_shape_shape_3")
            .attr_int("axis", 0)
            .input_shape("data", 5) // Shape input (represents shape of a 5D tensor)
            .add_input("indices", ArgType::Shape(3)) // Shape(3) indices
            .output_shape("output", 1) // Initial output, will be updated
            .build();

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        let _config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        // Should output Shape(3) since indices are Shape(3)
        match &node.outputs[0].ty {
            ArgType::Shape(rank) => {
                assert_eq!(*rank, 3, "Expected Shape(3) output for Shape(3) indices");
            }
            other => panic!("Expected Shape(3) output, got {:?}", other),
        }
    }

    #[test]
    fn test_gather_update_outputs_shape_with_tensor_indices() {
        // Test gather from Shape with Tensor indices -> Shape output with computed rank
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_gather_shape_tensor")
            .attr_int("axis", 0)
            .input_shape("data", 4) // Shape input
            .input_tensor_i64("indices", 1, None) // 1D tensor indices
            .output_shape("output", 1) // Initial output, will be updated
            .build();

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        let _config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        // Should output Shape(1) for 1D tensor indices (indices_rank = 1)
        match &node.outputs[0].ty {
            ArgType::Shape(rank) => {
                assert_eq!(*rank, 1, "Expected Shape(1) output for 1D tensor indices");
            }
            other => panic!("Expected Shape(1) output, got {:?}", other),
        }
    }

    #[test]
    fn test_gather_scalar_input() {
        // Test gather with scalar input - output should remain scalar
        // This tests the Reshape(scalar, [-1]) -> Scalar optimization path
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_gather_scalar")
            .attr_int("axis", 0)
            .add_input("data", ArgType::ScalarNative(crate::ir::DType::I64)) // Scalar input
            .add_input("indices", ArgType::ScalarNative(crate::ir::DType::I64)) // Scalar indices
            .add_output(
                "output",
                ArgType::Tensor(TensorType::new(crate::ir::DType::I64, 1, None)),
            ) // Will be updated
            .build();

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        let _config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        // Should output scalar when input is scalar
        match &node.outputs[0].ty {
            ArgType::ScalarNative(dtype) => {
                assert_eq!(*dtype, crate::ir::DType::I64);
            }
            other => panic!("Expected Scalar output, got {:?}", other),
        }
    }

    #[test]
    fn test_gather_is_noop_scalar_input() {
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_gather_noop")
            .attr_int("axis", 0)
            .add_input("data", ArgType::ScalarNative(crate::ir::DType::I64))
            .add_input("indices", ArgType::ScalarNative(crate::ir::DType::I64))
            .add_output(
                "output",
                ArgType::Tensor(TensorType::new(crate::ir::DType::I64, 1, None)),
            )
            .build();

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        // Gathering from a scalar is always a no-op
        assert!(processor.is_noop(&node));
    }

    #[test]
    fn test_gather_is_not_noop_tensor_input() {
        let mut node = create_test_node(0, 3, false).build();
        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        // Gathering from a tensor is not a no-op
        assert!(!processor.is_noop(&node));
    }

    #[test]
    fn test_gather_infer_types_lifts_constant_rank0_tensor_index() {
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_gather_rank0_tensor")
            .attr_int("axis", 0)
            .input_tensor_f32("data", 3, None)
            .input_scalar_tensor_i64("indices", Some(2))
            .output_tensor_f32("output", 2, None)
            .build_with_graph_data(16);

        node.inputs[1].ty = ArgType::Tensor(TensorType::new(DType::I64, 0, None));

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();

        assert!(node.inputs[1].is_constant()); // index is indeed constant
        processor.infer_types(&mut node, 16, &prefs).unwrap();
        assert!(node.inputs[1].is_static()); // was converted to static, can be eliminated
    }

    #[test]
    fn test_gather_infer_types_lifts_constant_scalar_tensor_index() {
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_gather_scalar_tensor")
            .attr_int("axis", 0)
            .input_tensor_f32("data", 3, None)
            .input_scalar_tensor_i64("indices", Some(2))
            .output_tensor_f32("output", 2, None)
            .build_with_graph_data(16);

        node.inputs[1].ty = ArgType::ScalarTensor(DType::I64);

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();

        assert!(node.inputs[1].is_constant()); // index is indeed constant
        processor.infer_types(&mut node, 16, &prefs).unwrap();
        assert!(node.inputs[1].is_static()); // was converted to static, can be eliminated
    }

    #[test]
    fn test_gather_infer_types_lifts_constant_scalar_index() {
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_gather_scalar")
            .attr_int("axis", 0)
            .input_tensor_f32("data", 3, None)
            .input_scalar_tensor_i64("indices", Some(2))
            .output_tensor_f32("output", 2, None)
            .build_with_graph_data(16);

        node.inputs[1].ty = ArgType::ScalarNative(DType::I64);

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();

        assert!(node.inputs[1].is_constant()); // index is indeed constant
        processor.infer_types(&mut node, 16, &prefs).unwrap();
        assert!(node.inputs[1].is_static()); // was converted to static, can be eliminated
    }

    #[test]
    fn test_gather_infer_types_dynamic_index() {
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_gather_dynamic")
            .attr_int("axis", 0)
            .input_tensor_f32("data", 3, None)
            .input_scalar_tensor_i64("indices", None)
            .output_tensor_f32("output", 2, None)
            .build_with_graph_data(16);

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();

        assert!(!node.inputs[1].is_constant()); // index is not constant
        processor.infer_types(&mut node, 16, &prefs).unwrap();
        assert!(!node.inputs[1].is_static()); // was not converted to static, can't be eliminated
    }

    #[test]
    fn test_gather_infer_types_constant_2d_index() {
        // 2d indexing is left as is for now
        let mut node = TestNodeBuilder::new(NodeType::Gather, "test_gather_constant_2d")
            .attr_int("axis", 0)
            .input_tensor_f32("data", 3, None)
            .input_scalar_tensor_i64("indices", Some(2))
            .output_tensor_f32("output", 2, None)
            .build_with_graph_data(16);

        node.inputs[1].ty = ArgType::Tensor(TensorType::new(DType::I64, 2, None));

        let processor = GatherProcessor;
        let prefs = OutputPreferences::new();

        assert!(node.inputs[1].is_constant()); // index is constant
        processor.infer_types(&mut node, 16, &prefs).unwrap();
        assert!(!node.inputs[1].is_static()); // was not converted to static, won't be eliminated
    }

    // TODO: Add test for out-of-bounds indices - Per spec (opset 11+), out-of-bounds indices should raise error - Missing bounds checking test
    // TODO: Add test for negative indices - Opset 11+ supports negative indices for backwards indexing - Missing negative indices test
    // TODO: Add test for indices type validation - Indices must be int32 or int64 per spec - Missing type constraint test
    // TODO: Add test for static shape computation - When both data and indices have static shapes, output should compute static shape - Missing shape inference test
    // TODO: Add test for zero-size tensors - Edge case where data or indices have 0 elements - Missing edge case test
    // TODO: Add test for different data types - Spec supports many types (all numeric, bool, string, complex) - Missing type coverage
    // TODO: Add test for unexpected attributes - Should reject unknown attributes per implementation - Missing attribute validation test
}