jw3126

# This reproduces a pytorch_lightning issue
# where gradient checkpointing + ddp results in nan loss
#
# * Run with gpus=1 and it works fine.
# * Run with gpus=4 and it loss becomes nan quickly
#
# See also https://forums.pytorchlightning.ai/t/gradient-checkpointing-ddp-nan/398
import torch
from torch import nn
from torch.nn import functional as F

jw3126

using Unitful
using Unitful: MeV, NoUnits, cm
using UnitfulRecipes

const h     = 6.626_070_040e-34*u"J*s"
const h_bar = h / (2pi)
const m_e   = 9.10938356e-31 * u"kg"
const c     = 299_792_458.0 * u"m/s"
const r_c   = h_bar / (c*m_e) # reduced compton wavelength of electron

jw3126

jw3126

# https://math.stackexchange.com/questions/1049788/haar-measure-of-an-angle-distance-ball-in-so3

using StatsBase
using LinearAlgebra
using StatPlots
using Plots
using Rotations

function sample_angles(N)
    map(1:N) do _