119 Commits

Author	SHA1	Message	Date
Matthew Johnson	ad4322c5da	playing around with flattening functions	2019-01-06 12:49:35 -08:00
Matthew Johnson	8f1bc997ca	add value-and-grad fun (closes #149)	2018-12-20 10:09:34 -08:00
Dougal Maclaurin	87922fdf13	Generalized make_jaxpr to handle python containers	2018-12-19 10:59:13 -05:00
Roy Frostig	e0ad5bb394	add TODO comment for containers in api.make_jaxpr	2018-12-17 11:42:45 -08:00
Roy Frostig	f4a8e03ce1	add a basic make_jaxpr transformation to the api module	2018-12-16 13:26:02 -08:00
Roy Frostig	3b8fdb050a	wrap "jit" around generated function name	2018-12-16 13:24:20 -08:00
Roy Frostig	b318f56928	generate name, module, and doctring for functions output from jit.	2018-12-16 13:04:29 -08:00
Dougal Maclaurin	9b835a180f	Added some docstrings to top-level transformations	2018-12-12 09:47:49 -05:00
Matthew Johnson	89349e5e6d	fix transpose issue in jacfwd and jacrev	2018-12-11 16:24:20 -08:00
Matthew Johnson	cbd9ace425	change vmap api to be curried (closes #78)	2018-12-11 12:52:09 -08:00
Matthew Johnson	b907bcd5ae	improve jaxlib wheel building script	2018-12-09 08:59:08 -08:00
Dougal Maclaurin	d89c7c2457	Merge branch 'master' into quickish-check	2018-12-07 11:13:41 -05:00
Matthew Johnson	50624bd978	rename in_bdims, out_bdims --> in_axes, out_axes	2018-12-07 06:53:29 -08:00
Dougal Maclaurin	2e4ff400fa	Fixed bug in vjp with constant-zero tangent outputs	2018-12-05 15:55:01 -05:00
Dougal Maclaurin	307d195577	Wrapped static args to jit to be hashed on id. This is conservative but simple and cheap.	2018-12-05 15:55:01 -05:00
Matthew Johnson	9101d66b4e	source sync ... PiperOrigin-RevId: 222487460	2018-11-21 20:32:33 -08:00
Matthew Johnson	3b3490f406	source sync ... PiperOrigin-RevId: 222483357	2018-11-21 20:32:33 -08:00
Peter Hawkins	5e60639bc5	source sync ... PiperOrigin-RevId: 222452709	2018-11-21 20:22:54 -08:00
Matthew Johnson	a30e858e59	populating source tree	2018-11-17 18:03:33 -08:00