#[burn_tensor_testgen::testgen(checkpoint)]
mod tests {
    use super::*;
    use burn_tensor::{Bool, Tensor, TensorData};

    #[test]
    fn test_autodiff_checkpoint_complicated_computation() {
        let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]);
        let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]);
        let data_2 = TensorData::from([[0.2, 7.0], [7.0, 7.0]]);
        let data_3 = TensorData::from([[0.3, 7.0], [7.0, 7.0]]);
        let data_4 = TensorData::from([[0.4, 7.0], [7.0, 7.0]]);

        let device = Default::default();
        let tensor_0 = TestAutodiffTensor::<2>::from_data(data_0, &device).require_grad();
        let tensor_1 = TestAutodiffTensor::from_data(data_1, &device).require_grad();
        let tensor_2 = TestAutodiffTensor::from_data(data_2, &device).require_grad();
        let tensor_3 = TestAutodiffTensor::from_data(data_3, &device).require_grad();
        let tensor_4 = TestAutodiffTensor::from_data(data_4, &device).require_grad();

        let tensor_5 = compute_bound_eager(tensor_0, tensor_1);
        let tensor_6 = compute_bound_lazy(tensor_2, tensor_3.clone());
        let tensor_7 = memory_bound_eager(tensor_3, tensor_4);
        let tensor_8 = compute_bound_lazy(tensor_6, tensor_7.clone());
        let tensor_9 = memory_bound_eager_scalar(tensor_7, 11.);
        let tensor_10 = memory_bound_lazy(tensor_5, tensor_8.clone());
        let tensor_11 = memory_bound_lazy(tensor_8, tensor_9);
        let tensor_12 = compute_bound_lazy(tensor_10, tensor_11);

        assert_checkpoint(tensor_12);
    }

    #[test]
    fn test_autodiff_checkpoint_with_missing_requirement() {
        let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]);
        let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]);

        let device = Default::default();
        let tensor_0 = TestAutodiffTensor::<2>::from_data(data_0, &device).require_grad();
        let tensor_1 = TestAutodiffTensor::from_data(data_1, &device); // does not require_grad

        let tensor_2 = memory_bound_eager(tensor_0, tensor_1);
        let tensor_3 = memory_bound_eager_scalar(tensor_2.clone(), 11.);
        let tensor_4 = memory_bound_eager_scalar(tensor_2.clone(), 11.);
        let tensor_5 = compute_bound_lazy(tensor_3, tensor_4);
        let tensor_6 = compute_bound_eager_scalar(tensor_5.clone(), 11.);
        let tensor_7 = memory_bound_eager(tensor_5, tensor_2);
        let tensor_8 = memory_bound_eager(tensor_6, tensor_7);

        assert_checkpoint(tensor_8);
    }

    #[test]
    fn test_autodiff_checkpoint_with_many_duplicates() {
        let data_0 = TensorData::from([[4.0, 7.0], [7.0, 7.0]]);

        let device = Default::default();
        let tensor_0 = TestAutodiffTensor::<2>::from_data(data_0, &device).require_grad();

        let tensor_1 = memory_bound_eager(tensor_0.clone(), tensor_0.clone());
        let tensor_2 = compute_bound_eager(tensor_0.clone(), tensor_0.clone());
        let tensor_3 = memory_bound_lazy(tensor_0.clone(), tensor_0.clone());
        let tensor_4 = compute_bound_lazy(tensor_0.clone(), tensor_0.clone());

        let tensor_5 = memory_bound_eager(tensor_1.clone(), tensor_0.clone());
        let tensor_6 = memory_bound_eager(tensor_0.clone(), tensor_5.clone());
        let tensor_7 = compute_bound_lazy(tensor_3.clone(), tensor_5.clone());
        let tensor_8 = compute_bound_eager(tensor_4.clone(), tensor_2.clone());
        let tensor_9 = memory_bound_lazy(tensor_6, tensor_7);
        let tensor_10 = memory_bound_eager(tensor_0, tensor_9);
        let tensor_11 = memory_bound_eager_scalar(tensor_10, 9.);
        let tensor_12 = compute_bound_lazy(tensor_8, tensor_11);

        assert_checkpoint(tensor_12);
    }

    #[test]
    fn test_autodiff_checkpoint_with_long_chain_of_eager_memory_bound() {
        let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]);
        let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]);
        let data_2 = TensorData::from([[0.2, 7.0], [7.0, 7.0]]);
        let data_3 = TensorData::from([[0.3, 7.0], [7.0, 7.0]]);
        let data_4 = TensorData::from([[0.4, 7.0], [7.0, 7.0]]);

        let device = Default::default();
        let tensor_0 = TestAutodiffTensor::<2>::from_data(data_0, &device).require_grad();
        let tensor_1 = TestAutodiffTensor::from_data(data_1, &device);
        let tensor_2 = TestAutodiffTensor::from_data(data_2, &device).require_grad();
        let tensor_3 = TestAutodiffTensor::from_data(data_3, &device).require_grad();
        let tensor_4 = TestAutodiffTensor::from_data(data_4, &device).require_grad();

        let tensor_5 = memory_bound_eager(tensor_0, tensor_1.clone());
        let tensor_6 = memory_bound_eager(tensor_5, tensor_2);
        let tensor_7 = memory_bound_eager(tensor_6, tensor_3);
        let tensor_8 = memory_bound_eager(tensor_7, tensor_4);
        let tensor_9 = memory_bound_lazy(tensor_8, tensor_1);

        assert_checkpoint(tensor_9)
    }

    #[test]
    fn test_autodiff_checkpoint_half_sub_graph_not_tracked() {
        let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]);
        let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]);
        let data_2 = TensorData::from([[0.2, 7.0], [7.0, 7.0]]);
        let data_3 = TensorData::from([[0.3, 7.0], [7.0, 7.0]]);
        let data_4 = TensorData::from([[0.4, 7.0], [7.0, 7.0]]);
        let data_5 = TensorData::from([[0.5, 7.0], [7.0, 7.0]]);

        let device = Default::default();
        let tensor_0 = TestAutodiffTensor::<2>::from_data(data_0, &device);
        let tensor_1 = TestAutodiffTensor::from_data(data_1, &device);
        let tensor_2 = TestAutodiffTensor::from_data(data_2, &device);
        let tensor_3 = TestAutodiffTensor::from_data(data_3, &device).require_grad();
        let tensor_4 = TestAutodiffTensor::from_data(data_4, &device).require_grad();
        let tensor_5 = TestAutodiffTensor::from_data(data_5, &device).require_grad();

        let tensor_6 = memory_bound_lazy(tensor_0, tensor_1);
        let tensor_7 = compute_bound_eager(tensor_6, tensor_2);

        let tensor_8 = memory_bound_eager(tensor_3, tensor_4);
        let tensor_9 = compute_bound_lazy(tensor_8, tensor_5);

        let tensor_10 = compute_bound_lazy(tensor_7, tensor_9);

        assert_checkpoint(tensor_10);
    }

    #[test]
    fn test_autodiff_checkpoint_very_complex() {
        let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]);
        let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]);
        let data_2 = TensorData::from([[0.2, 7.0], [7.0, 7.0]]);
        let data_3 = TensorData::from([[0.3, 7.0], [7.0, 7.0]]);
        let data_4 = TensorData::from([[0.4, 7.0], [7.0, 7.0]]);

        let device = Default::default();
        let tensor_0 = TestAutodiffTensor::<2>::from_data(data_0, &device).require_grad();
        let tensor_1 = TestAutodiffTensor::from_data(data_1, &device);
        let tensor_2 = TestAutodiffTensor::from_data(data_2, &device).require_grad();
        let tensor_3 = TestAutodiffTensor::from_data(data_3, &device).require_grad();
        let tensor_4 = TestAutodiffTensor::from_data(data_4, &device).require_grad();

        let tensor_5 = memory_bound_eager_scalar(tensor_0, 8.);
        let tensor_6 = memory_bound_lazy(tensor_5.clone(), tensor_1.clone());
        let tensor_7 = compute_bound_lazy(tensor_6.clone(), tensor_6);
        let tensor_8 = memory_bound_lazy(tensor_1.clone(), tensor_5.clone());
        let tensor_9 = memory_bound_eager_scalar(tensor_7.clone(), 7.);
        let tensor_10 = compute_bound_eager(tensor_5, tensor_8);
        let tensor_11 = memory_bound_eager(tensor_2.clone(), tensor_9);
        let tensor_12 = memory_bound_lazy(tensor_2.clone(), tensor_2);
        let tensor_13 = compute_bound_eager(tensor_10.clone(), tensor_11);
        let tensor_14 = compute_bound_eager_scalar(tensor_3, 8.);
        let tensor_15 = compute_bound_lazy(tensor_4, tensor_12);
        let tensor_16 = memory_bound_lazy(tensor_10, tensor_7);
        let tensor_17 = compute_bound_lazy(tensor_13, tensor_1);
        let tensor_18 = memory_bound_eager(tensor_15, tensor_16);
        let tensor_19 = compute_bound_eager(tensor_14, tensor_17);
        let tensor_20 = memory_bound_lazy(tensor_18, tensor_19);
        let tensor_21 = memory_bound_eager_scalar(tensor_20, 8.);

        assert_checkpoint(tensor_21)
    }

    fn assert_checkpoint<const D: usize>(tensor: TestAutodiffTensor<D>) {
        // Assert is not explicit here, but the test can fail
        // - when a tensor is actually required more than n_required, it won't be found and will panic
        // - when a tensor is actually required less than n_required, the backward states map won't be
        //   empty and will fail the assertion within the backward code, same for retro_forwards
        tensor.backward();
    }

    // Does not save its state and does not need its parents
    fn memory_bound_eager<const D: usize>(
        tensor_a: TestAutodiffTensor<D>,
        tensor_b: TestAutodiffTensor<D>,
    ) -> TestAutodiffTensor<D> {
        tensor_a.add(tensor_b)
    }
    fn memory_bound_eager_scalar<const D: usize>(
        tensor_a: TestAutodiffTensor<D>,
        b: f32,
    ) -> TestAutodiffTensor<D> {
        tensor_a.add_scalar(b)
    }

    // Saves its own state and does not need its parents
    fn compute_bound_eager<const D: usize>(
        tensor_a: TestAutodiffTensor<D>,
        tensor_b: TestAutodiffTensor<D>,
    ) -> TestAutodiffTensor<D> {
        let mask =
            Tensor::<TestAutodiffBackend, D, Bool>::empty(tensor_a.shape(), &tensor_a.device());
        tensor_a.mask_where(mask, tensor_b)
    }
    fn compute_bound_eager_scalar<const D: usize>(
        tensor_a: TestAutodiffTensor<D>,
        b: f32,
    ) -> TestAutodiffTensor<D> {
        let mask =
            Tensor::<TestAutodiffBackend, D, Bool>::empty(tensor_a.shape(), &tensor_a.device());
        tensor_a.mask_fill(mask, b)
    }

    // Does not save its state and needs its parents
    fn memory_bound_lazy<const D: usize>(
        tensor_a: TestAutodiffTensor<D>,
        tensor_b: TestAutodiffTensor<D>,
    ) -> TestAutodiffTensor<D> {
        tensor_a.mul(tensor_b)
    }

    // Saves its own state and needs its parents
    fn compute_bound_lazy<const D: usize>(
        tensor_a: TestAutodiffTensor<D>,
        tensor_b: TestAutodiffTensor<D>,
    ) -> TestAutodiffTensor<D> {
        tensor_a.matmul(tensor_b)
    }
}
