import pdb
import warnings

from jax import custom_vjp


@custom_vjp
def debug_identity(x):
    """
    acts as identity, but inserts a pdb trace on the backwards pass
    """
    warnings.warn('Using a module intended for debugging')
    return x


def _debug_fwd(x):
    warnings.warn('Using a module intended for debugging')
    return x, x


# noinspection PyUnusedLocal
def _debug_bwd(x, g):
    pdb.set_trace()
    return g


debug_identity.defvjp(_debug_fwd, _debug_bwd)

[((12, 4, 12, 58), 'warnings.warn', 'warnings.warn', ({(12, 18, 12, 57): '"""Using a module intended for debugging"""'}, {}), "('Using a module intended for debugging')", False, 'import warnings\n'), ((17, 4, 17, 58), 'warnings.warn', 'warnings.warn', ({(17, 18, 17, 57): '"""Using a module intended for debugging"""'}, {}), "('Using a module intended for debugging')", False, 'import warnings\n'), ((23, 4, 23, 19), 'pdb.set_trace', 'pdb.set_trace', ({}, {}), '()', False, 'import pdb\n')]