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)

[((180, 234), 'warnings.warn', 'warnings.warn', (['"""Using a module intended for debugging"""'], {}), "('Using a module intended for debugging')\n", (193, 234), False, 'import warnings\n'), ((273, 327), 'warnings.warn', 'warnings.warn', (['"""Using a module intended for debugging"""'], {}), "('Using a module intended for debugging')\n", (286, 327), False, 'import warnings\n'), ((401, 416), 'pdb.set_trace', 'pdb.set_trace', ([], {}), '()\n', (414, 416), False, 'import pdb\n')]