feat: add examples/dp_sgd_transformer.py for DP training of a transformer.

[Neerajpathak07]

Fixes:- #86

What is the purpose of this pull request?

This pull request:-

• Adds a single file to jax_privacy/examples to add a simple end-to-end and realistic example usage of the library.
• Utilizes core components like, optax, clipped_grad & noise_addition.
• References:- https://github.com/google-deepmind/jax_privacy/blob/main/examples/jax_api_example.py.

[ryan112358]

I really like what you've done here, thanks for the PR. A few more comments then I think will be good to submit!

[Neerajpathak07]

@ryan112358 Glad to hear that, a big thanks for speed running this PR by providing crucial reviews!!

[ryan112358]

This has passed the first round of approval, now you mainly need to fix some formatting issues that came up in our internal checks. Please run flake8, pylint, and pytype on the add file and fix any issues encountered. Here is the list of problems that we see internally

Line too long (84/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:19:1
Line too long (89/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:23:1
Line too long (98/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:24:1
Line too long (111/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:25:1
Line too long (89/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:26:1
Line too long (115/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:27:1
Line too long (104/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:28:1
Line too long (119/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:29:1
Line too long (83/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:30:1
Line too long (83/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:33:1
Line too long (98/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:54:1
Line too long (81/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:74:1
Line too long (81/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:75:1
Line too long (82/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:93:1
Line too long (82/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:94:1
Line too long (82/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:95:1
Line too long (81/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:99:1
Line too long (95/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:118:1
Line too long (88/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:124:1
Line too long (87/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:149:1
Line too long (90/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:182:1
Line too long (90/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:267:1
Line too long (95/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:270:1
Line too long (93/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:282:1
Bad indentation. Found 8 spaces, expected 6 [bad-indentation]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:301:1
Line too long (83/80) [line-too-long]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:305:1
Quote delimiter " is inconsistent with the rest of the file [g-inconsistent-quotes]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:185:1
Quote delimiter " is inconsistent with the rest of the file [g-inconsistent-quotes]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:301:1
Quote delimiter " is inconsistent with the rest of the file [g-inconsistent-quotes]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:304:1
Quote delimiter " is inconsistent with the rest of the file [g-inconsistent-quotes]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:306:1
Quote delimiter " is inconsistent with the rest of the file [g-inconsistent-quotes]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:307:1
Quote delimiter " is inconsistent with the rest of the file [g-inconsistent-quotes]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:310:1
Unused argument 'num_heads' [unused-argument]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:54:59
Unused variable 'i' [unused-variable]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:64:7
Possible unbalanced tuple unpacking with sequence defined at line 275: left side has 3 labels, right side has 4 values [unbalanced-tuple-unpacking]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:292:5
Unused numpy imported as np [unused-import]
//depot/google3/third_party/py/jax_privacy/examples/dp_sgd_transformer.py:48:1
Invalid import order.
Diff:
from absl import app
import dp_accounting
import jax
from jax import random
import jax.numpy as jnp
import jax_privacy
from jax_privacy import batch_selection

• from jax_privacy.experimental import execution_plan
  from jax_privacy import noise_addition
  from jax_privacy.accounting import accountants
  from jax_privacy.accounting import analysis
  from jax_privacy.accounting import calibrate

• from jax_privacy.experimental import execution_plan
  import numpy as np
  import optax
  import tensorflow as tf
  import tensorflow_datasets as tfds
  Copy and pastable ordering:
  from absl import app
  import dp_accounting
  import jax
  from jax import random
  import jax.numpy as jnp
  import jax_privacy
  from jax_privacy import batch_selection
  from jax_privacy import noise_addition
  from jax_privacy.accounting import accountants
  from jax_privacy.accounting import analysis
  from jax_privacy.accounting import calibrate
  from jax_privacy.experimental import execution_plan
  import numpy as np
  import optax
  import tensorflow as tf
  import tensorflow_datasets as tfds
  [g-bad-import-order]

[RamSaw]

Thank you very much for your contribution! I agree with Ryan that this example is a very valuable one!

I did my review and left few comments. Ryan or I might add two comments soon once I clarify some things for myself.

[RamSaw]

There are linter errors, could you sync your PR to head, then you will see them. We enabled linter checks for examples dir few minutes go :)

[Neerajpathak07]

@RamSaw CI is running clean now just the comment on import tensoflow_datasets as tfds needs to be addressed then I think this should cross the Finish line. Feel free to LMK if there is anything else.

[RamSaw]

There are some complains from internal linter, I will figure it out too soon.

I need to merge #133, then it should catch the formatting issues in your newly added file and give one command to apply the correct formatting. But for some reason it wants to format already existing, so I have to figure out what is wrong.

[Neerajpathak07]

@RamSaw awaiting #133 to be merged after which I'll pull from it as of now the import error has been resolved Thanks for the suggestion!!